#!/usr/bin/env python3

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'

class Client:
    def __init__(self, fingerprint: bytes, common_name: str, admin: bool = False):
        self.fingerprint = fingerprint
        self.common_name = common_name
        self.admin = admin

def get_client_info(client_fingerprint: bytes):
    with dbm.open('server/data/conf/client_fingerprints') as db:
        if client_fingerprint not in db.keys():
            raise KeyError('Client not registered')
        common_name = db[client_fingerprint].decode()
        db.close()
    with open('server/data/conf/client_admin_rights') as carf:
        admins = carf.read().split('\n')
        carf.close()
    return Client(fingerprint = client_fingerprint, common_name = common_name, admin = client_fingerprint.hex() in admins)


def process_command(command, client: Client):
    print(client.fingerprint, client.common_name, client.admin)
    response = b''
    if command[0] == 0x11:
        if not client.admin:
            response = b'\xb7'
    return response



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/data/conf/client_fingerprints', 'c') as db:
        if key_hash not in db.keys():
            await sendmsg('ERR', writer, aesgcm, server_nonce)
            raise InvalidCertificateError('client not known')
        if not validate_cert(cert, db[key_hash].decode(), CLIENT_CERT_ISSUER_NAME):
            await sendmsg('ERR', writer, aesgcm, server_nonce)
            raise InvalidCertificateError('client certificate not trusted')

    await sendmsg('SCD', writer, aesgcm, server_nonce)
    print('CLIENT AUTHENTICATED')

    while True:
        try:
            await sendmsg(process_command(await readmsg(reader, aesgcm, client_nonce), get_client_info(hashlib.sha256(cert.public_key().public_bytes(encoding=serialization.Encoding.DER, format=serialization.PublicFormat.SubjectPublicKeyInfo)).digest())), writer, aesgcm, server_nonce)
        except Exception as e:
            print(f'{str(type(e))[8:-2]}: {e}')
            break

    writer.close()
    await writer.wait_closed()
    print(f'Connection to {addr} 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())