import socket
import threading
import sys
import hashlib
import os
import OpenSSL
import Encryptor
import time
from OpenSSL import SSL
from OpenSSL import crypto
from Crypto import Random
from Crypto.Cipher import AES
from diffiehellman.diffiehellman import DiffieHellman
from OpenCA import createCSR

BUFSIZE=32768

class Client:
    sock= socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    mujAES=None
    cli=None
    cislo=0
    def __init__(self,adress):

        self.sock.connect((adress, 9876))
        print("Připojen k serveru na adrese" + adress)
        #zjisti zda už u sebe má certifikát nebo alespoň certreq, pokud ano pošle ho serveru, pokud ne vytvoří a pošle žádost
        if os.path.isfile('User.private.pem') == False or os.path.isfile('client1.cert.pem') == False:
            createCSR('User','heslo',{'CN':'USER_FQDN'})
            self.poslatCertReq()
            print("Certifikát nenalezen, žádám o nový.")

        else:
            print("Certifikát nalezen, odesílám autoritě.")
            self.cert=OpenSSL.crypto.load_certificate(crypto.FILETYPE_PEM,open('client1.cert.pem').read())
            self.poslatCert()

       #diffie hellman generace veřejných klíčů
        self.cli=DiffieHellman()
        self.cli.generate_public_key()

        self.vyzadatKlic()

        iThread=threading.Thread(target=self.sendMessage)
        iThread.daemon=True
        iThread.start()
        self.run()


    def run(self):

        while True:
            data=self.sock.recv(BUFSIZE)
            #zde se resi typy zprav
            if not data:
                #ukonceni komunikace
                break

            elif data[0:1]==b"\x11":
                #kdyz prijde ridici znak x11-posleme na vyzadani klic
                self.nastavitKlic(data[1:])
                self.poslatKlic()
                self.cislo=1

            elif data[0:1]==b"\x12":
                #kdyz prijde ridici znak x12 tak si nastavime klic ktery nasleduje po tomto bytu
                self.nastavitKlic(data[1:])
                self.cislo=2

            #elif data[0:1]==b"\x13":
                #nastaveni klice v pripade ze byl vyzadan
             #   self.nastavitKlic(data[1:])

            elif data[0:1]==b"\x21":
                #kdyz prijde ridici znak x21 - prisel hotovy certifikat, nastavime ho
                print("Certifikát přijat.")
                with open('client1.cert.pem','wb') as cert:
                    cert.write(data[1:])
                self.nastavitCert(data[1:])

            elif data[0:1]==b"\x22":
                #ridici znak x22 - server nas pozadal o novou zadost o certifikat
                createCSR('User','heslo',{'CN':'USER_FQDN'})
                self.poslatCertReq()

            elif data[0:1]==b"\x98":
                #nezasifrovana komunikace - nemelo by nastat
                print(data.decode())

            elif data[0:1]==b"\x99":
                #sifrovana komunikace
                data=self.mujAES.decrypt(data[1:])
                try:
                    print("Klient "+str(self.cislo)+":"+data.decode())
                except:
                    continue

    def nastavitCert(self,data):
        #funkce která uloží certifikát a vypíše ho
        self.cert=OpenSSL.crypto.load_certificate(crypto.FILETYPE_PEM,data)
        print("Certifikát uložen.")
        cert=OpenSSL.crypto.load_certificate(crypto.FILETYPE_PEM,open('client1.cert.pem').read())
        print(OpenSSL.crypto.dump_certificate(crypto.FILETYPE_TEXT,cert).decode())

    def poslatCert(self):
        #funkce která posílá existující certifikát serveru pro ověření
        print("Posílám certifikát")
        text=OpenSSL.crypto.dump_certificate(crypto.FILETYPE_PEM,self.cert)
        self.sock.send(b'\x23'+text)

    def poslatCertReq(self):
        #posle certifikat na podepsani¨
        with open('User.CSR.pem') as cert:
            certificate = OpenSSL.crypto.load_certificate_request(crypto.FILETYPE_PEM, cert.read())
            certext = OpenSSL.crypto.dump_certificate_request(crypto.FILETYPE_PEM, certificate)
            print("Žádost" + str(certext))
            self.sock.send(b'\x25'+certext)
            print('Žádost odeslána.')

    def poslatKlic(self):
        #posle ridici znak nasledovany klicem
        self.sock.send(b'\x12'+str(self.cli.public_key).encode())
        print("Posílám veřejný klíč.")

    def nastavitKlic(self,data):
        #nastavuje klic na zaklade dat ktere dostane
        self.cli.generate_shared_secret(int(data.decode()),echo_return_key=True)
        superklic=str(self.cli.shared_secret)
        klic=hashlib.sha256(superklic.encode()).hexdigest()[:32]
        print("Tajný klíč vytvořen.")
        #self.cislo=1
        print("Klíč pro AES:" +klic)
        self.mujAES=Encryptor.Encryptor(klic)
        #self.sock.send(b'\x13'+str(self.cli.public_key).encode())

#    def jinenastaveniklice(self,data):
 #       #nastavuje klic v případě že jsme ho vyžádali
  #      self.cli.generate_shared_secret(int(data.decode()),echo_return_key=True)
   #     superklic=str(self.cli.shared_secret)
    #    klic=hashlib.sha256(superklic.encode()).hexdigest()[:32]
     #   print("Tajný klíč vytvořen.")
      #  self.cislo=2
       # print("Klíč pro AES:" + klic)
        #self.mujAES=Encryptor.Encryptor(klic)

    def vyzadatKlic(self):
        #nemam klic ale chci, poslu ridici znak
        print("Žádám o veřejný klíč pro DH.")
        self.sock.send(b'\x11'+str(self.cli.public_key).encode())

    def sendMessage(self):
        #hlavni chatova smycka

        while True:
            msg=str(input(""))

            if self.mujAES is not None:
                msg=self.mujAES.encrypt(msg.encode())
                self.sock.send(b'\x99' + msg)
            else:
                msg=msg.encode()
                self.sock.send(b'\x98'+msg)


adresa='127.0.0.1'
client=Client(adresa)