import os
import pickle
import string
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

CURVE = ec.SECP256R1()

def GENERATE_DH():
    return ec.generate_private_key(CURVE)

def DH(dh_pair, dh_pub):
    return dh_pair.exchange(ec.ECDH(), dh_pub)

class MessengerServer:
    def __init__(self, server_signing_key, server_decryption_key):
        self.server_signing_key = server_signing_key
        self.server_decryption_key = server_decryption_key

    def decryptReport(self, ct):
        raise Exception("not implemented!")
        return

    def signCert(self, cert):
        return self.server_signing_key.sign(cert, ec.ECDSA(hashes.SHA256()))

class MessengerClient:

    def __init__(self, name, server_signing_pk, server_encryption_pk):
        self.name = name
        self.server_signing_pk = server_signing_pk
        self.server_encryption_pk = server_encryption_pk
        self.conns = {}
        self.certs = {}
        self.dh = GENERATE_DH()
        self.certificate = None

    def generateCertificate(self):
        self.certificate = pickle.dumps((self.name, self.dh.public_key().public_bytes(serialization.Encoding.X962,
        serialization.PublicFormat.UncompressedPoint)))
        return self.certificate

    def receiveCertificate(self, certificate, signature):
        self.server_signing_pk.verify(signature, certificate, ec.ECDSA(hashes.SHA256()))
        name, public_bytes = pickle.loads(certificate)
        public_key = ec.EllipticCurvePublicKey.from_encoded_point(CURVE, public_bytes)
        self.certs[name] = public_key

    def sendMessage(self, name, message):
        raise Exception("not implemented!")
        return

    def receiveMessage(self, name, header, ciphertext):
        raise Exception("not implemented!")
        return

    def report(self, name, message):
        report = (name, message)
        plaintext = pickle.dumps(report)
        temp_key = GENERATE_DH()
        temp_bytes = temp_key.public_key().public_bytes(
            serialization.Encoding.X962,
            serialization.PublicFormat.UncompressedPoint,
        )
        shared_secret = DH(temp_key, self.server_encryption_pk)
        hkdf = HKDF(
            algorithm=hashes.SHA256(),
            length=32,
            salt=None,
            info=b""
        )
        key = hkdf.derive(shared_secret)
        nonce = os.urandom(12)
        aesgcm = AESGCM(key)
        ciphertext = aesgcm.encrypt(nonce, plaintext, temp_bytes) # not sure if passing in temp_bytes is correct, but it seems like it should be
        data = {
            "bytes": temp_bytes,
            "nonce": nonce,
            "ciphertext": ciphertext
        }
        return plaintext.decode("latin1"), pickle.dumps(data)