# Plik pomocniczy
import os, base64, json
from Crypto.Cipher import PKCS1_OAEP, AES
from Crypto.PublicKey import RSA, ECC
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)
    return message

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

# 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")

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

    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

# 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

    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

    return 1

        

