import sys
import os
sys.path.append(os.getcwd())
import asyncio
from lib.crypto_utils import (
    generate_keypair,
    serialize_public_key,
    deserialize_public_key,
    derive_aes_key,
)
from lib.jebp_utils import sendmsg, readmsg, validate_cert, InvalidCertificateError
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.backends import default_backend
from cryptography import x509
import base64
import dbm
import hashlib


HOST = "0.0.0.0"
PORT = 8888
VERSION = 'jebp 1.0'
CERT_FILE = 'server/sec/server.crt.pem'
CLIENT_CERT_ISSUER_NAME = 'jCloudCA-Root-CA'

async def handle_client(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
    addr = writer.get_extra_info('peername')
    print(f'Connected to {addr}')


    await sendmsg(VERSION, writer)
    
    # BEGIN ENCRYPTION HANDSHAKE

    server_priv, server_pub = generate_keypair()

    client_pub_bytes = await readmsg(reader)
    client_pub = deserialize_public_key(client_pub_bytes)

    await sendmsg(serialize_public_key(server_pub), writer)

    aes_key = derive_aes_key(server_priv, client_pub)
    aesgcm = AESGCM(aes_key)

    client_nonce = await readmsg(reader)

    server_nonce = os.urandom(12)
    await sendmsg(server_nonce, writer)

    await sendmsg(await readmsg(reader, aesgcm, client_nonce), writer)


    # BEGIN SERVER AUTHENTICATION HANDSHAKE


    await readmsg(reader)
    with open(CERT_FILE, 'rb') as certfile:
        cert_data = certfile.read()
        certfile.close()

    await sendmsg(base64.b64decode(cert_data.replace(b'-----BEGIN CERTIFICATE-----', b'').replace(b'-----END CERTIFICATE-----', b'').strip()), writer, aesgcm, server_nonce)


    # BEGIN CLIENT AUTHENTICATION HANDSHAKE

    cert = x509.load_der_x509_certificate(await readmsg(reader, aesgcm, client_nonce))
    key_hash = hashlib.sha256(cert.public_key().public_bytes(encoding=serialization.Encoding.DER, format=serialization.PublicFormat.SubjectPublicKeyInfo)).digest()
    with dbm.open('server/config/clients/fingerprints', 'c') as db:
        if key_hash not in db.keys():
            raise InvalidCertificateError('client not known')
        if not validate_cert(cert, db[key_hash].decode(), CLIENT_CERT_ISSUER_NAME):
            raise InvalidCertificateError('client certificate not trusted')

    print('Client authenticated')

    writer.close()
    await writer.wait_closed()


async def main():
    server = await asyncio.start_server(handle_client, HOST, PORT)
    addr = ', '.join(str(sock.getsockname()) for sock in server.sockets)
    print(f'Listening on {addr}')

    async with server:
        await server.serve_forever()


if __name__ == '__main__':
    asyncio.run(main())