# Plik pomocniczy
import os, base64, json
from Crypto.Cipher import PKCS1_OAEP, AES
from Crypto.PublicKey import RSA, ECC
import r09_diffiehellman as dh
from binascii import hexlify, unhexlify
from base64 import b64encode, b64decode

# Metody szyfrujące
def encrypt(message, usePKI, useDH, dhSecret):
    if usePKI == True:
        message = encrypt_rsa(message)
    if useDH == True:
        message = encrypt_dh(message, dhSecret)
    return message

# Metody odszyfrowujące
def decrypt(message, usePKI, useDH, dhSecret):
    if useDH == True:
        message = decrypt_dh(message, dhSecret)
    if usePKI == True:
        message = decrypt_rsa(message)
    return message

# Usuwanie wygenerowanych certyfikatów dh
def remove_dh_certs():
    try:
        os.remove("klucz_prywatny_dh_klient.pem")
        os.remove("klucz_publiczny_dh_klient.pem")
        os.remove("klucz_prywatny_dh_serwer.pem")
        os.remove("klucz_publiczny_dh_serwer.pem")
    except:
        return 0

# Usuwanie wygenerowanych certyfikatów RSA
def remove_rsa_certs():
    try:
        os.remove("klucz_prywatny_klient.pem")
        os.remove("klucz_publiczny_klient.pem")
    except:
        return 0

# Generowanie certyfikatów Diffiego-Hellmana dla klienta
def gen_client_DH():
    clientDH =  dh.DiffieHellman(2,17,1024)
    
    privateKey = str(clientDH.privateKey).encode()
    fd = open("klucz_prywatny_dh_klient.pem", "wb")
    fd.write(privateKey)
    fd.close()

    publicKey = str(clientDH.publicKey).encode()
    fd = open("klucz_publiczny_dh_klient.pem", "wb")
    fd.write(publicKey)
    fd.close()

    clientDHSet = clientDH
    return clientDH

# Generowanie certyfikatów Diffiego-Hellmana dla serwera
def gen_server_DH():

    svrDH =  dh.DiffieHellman(2,17,1024)

    privateKey = str(svrDH.privateKey).encode()
    fd = open("klucz_prywatny_dh_serwer.pem", "wb")
    fd.write(privateKey)
    fd.close()

    publicKey = str(svrDH.publicKey).encode()
    fd = open("klucz_publiczny_dh_serwer.pem", "wb")
    fd.write(publicKey)
    fd.close()

    key = (open('klucz_publiczny_dh_serwer.pem').read())

    serverDHSet = svrDH
    return svrDH

# Generowanie certyfikatów RSA
def gen_rsa_certs():
    # r08_RSA_certyfikaty.py
    
    # Generowanie pary klucz publiczny – klucz prywatny; długość klucza to 4096 bitów (512 bajtów)

    new_key = RSA.generate(4096, e=65537)

    # Klucz prywatny w formacie PEM
    private_key = new_key.exportKey("PEM")

    # Klucz publiczny w formacie PEM
    public_key = new_key.publickey().exportKey("PEM")

    # print(private_key)
    fd = open("klucz_prywatny_klient.pem", "wb")
    fd.write(private_key)
    fd.close()

    # print(public_key)
    fd = open("klucz_publiczny_klient.pem", "wb")
    fd.write(public_key)
    fd.close()

# Odszyfrowywanie RSA
def decrypt_rsa(ciphertext):
    key = RSA.importKey(open('klucz_prywatny_klient.pem').read())
    cipher = PKCS1_OAEP.new(key)
    plaintext = cipher.decrypt(ciphertext)
    return plaintext

# Szyfrowanie RSA
def encrypt_rsa(message):
    key = RSA.importKey(open('klucz_publiczny_klient.pem').read())
    cipher = PKCS1_OAEP.new(key)
    ciphertext = cipher.encrypt(message)

    plaintext = decrypt_rsa(ciphertext)

    return ciphertext

# Szyfrowanie za pomocą Diffie-Hellman - ECC
def encrypt_dh(plaintext, dhSecret):
    # Szyfrowanie za pomocą współdzielonego klucza uzyskanego z klucza klienta (klucz prywatny) i serwera (klucz publiczny)
    ciphertext = encrypt_AES_GCM(plaintext,dhSecret)
    ciphertext = ciphertext.encode()
    # reverse = decrypt_AES_GCM(ciphertext, dhSecret)
    return ciphertext

# Odszyfrowywanie za pomocą Diffie-Hellman - ECC
def decrypt_dh(ciphertext, dhSecret):
    # Odszyfrowywanie za pomocą współdzielonego klucza uzyskanego z klucza klienta (klucz prywatny) i serwera (klucz publiczny)
    ciphertext = ciphertext.decode('utf-8')
    plaintext = decrypt_AES_GCM(ciphertext,dhSecret)
    # reverse = encrypt_AES_GCM(ciphertext, dhSecret)
    return plaintext

# Generowanie certyfikatów ECC
def gen_ecc_certs():
    key = ECC.generate(curve='P-256')
    f = open('moj_klucz_prywatny.pem','wt')
    f.write(key.export_key(format='PEM'))
    f.close()
    f = open('moj_klucz_prywatny.pem','rt')
    key = ECC.import_key(f.read())
    print(key)

# Szyfrowanie AES-GCM
def encrypt_AES_GCM(msg, secretKey):
    aesCipher = AES.new(secretKey, AES.MODE_GCM)
    ct, authTag = aesCipher.encrypt_and_digest(msg)

    ct = hexlify(ct)
    ct = ct.decode('utf-8')
    authTag = hexlify(authTag)
    authTag = authTag.decode('utf-8')
    noncea = hexlify(aesCipher.nonce)
    nonce = noncea.decode('utf-8')
    
    ciphertext = json.dumps({'nonce':nonce, 'ciphertext':ct, 'tag':authTag})

    return ciphertext

# Odszyfrowywanie AES-GCM
def decrypt_AES_GCM(encryptedMsg, secretKey):
   
    b64 = json.loads(encryptedMsg)
    
    nonce = str(b64['nonce'])
    nonce = nonce.encode()
    nonce = unhexlify(nonce)
    
    ct = str(b64['ciphertext'])
    ct = ct.encode()
    ct = unhexlify(ct)

    authTag = str(b64['tag'])
    authTag = authTag.encode()
    authTag = unhexlify(authTag)
    
    aesCipher = AES.new(secretKey, AES.MODE_GCM, nonce)
    nonce = b64encode(aesCipher.nonce).decode('utf-8')
    plaintext = aesCipher.decrypt_and_verify(ct, authTag)
    return plaintext
   
# Szyfrowanie AES-CTR
def encrypt_AES_CTR(msg, secretKey):
    cipher = AES.new(secretKey, AES.MODE_CTR)
    ct_bytes = cipher.encrypt(msg)
    nonce = b64encode(cipher.nonce).decode('utf-8')
    ct = b64encode(ct_bytes).decode('utf-8')
    ciphertext = json.dumps({'nonce':nonce, 'ciphertext':ct})
    return ciphertext

# Odszyfrowywanie AES-CTR
def decrypt_AES_CTR(msg, secretKey):
    b64 = json.loads(msg)
    nonce = b64decode(b64['nonce'])
    ct = b64decode(b64['ciphertext'])
    cipher = AES.new(secretKey, AES.MODE_CTR, nonce=nonce)
    plaintext = cipher.decrypt(ct)
    return plaintext

# Sprawdzanie poleceń po stronie klienta
def check_client_command(data):
    if data == b'addPKI':
        gen_rsa_certs()
        return 11
    elif data == b'removePKI':
        usePKI = False
        return 10
    elif data == b'addDH':
        return 21
    elif data == b'removeDH':
        usePKI = False
        return 20
    return 1

# Sprawdzanie poleceń po stronie serwera
def check_server_command(data):
    if data == b'addPKI':
        return 11
    if data == b'removePKI':
        useDH = False
        return 10
    if data == b'addDH':
        return 21
    if data == b'removeDH':
        useDH = False
        return 20
    return 1

        

