import argparse
import os

import cryptography.hazmat.primitives.serialization as serialization
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import dh
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives.serialization import (
    Encoding,
    NoEncryption,
    PrivateFormat,
    PublicFormat,
)

p = 0xFFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD129024E088A67CC74020BBEA63B139B22514A08798E3404DDEF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245E485B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7EDEE386BFB5A899FA5AE9F24117C4B1FE649286651ECE45B3DC2007CB8A163BF0598DA48361C55D39A69163FA8FD24CF5F83655D23DCA3AD961C62F356208552BB9ED529077096966D670C354E4ABC9804F1746C08CA18217C32905E462E36CE3BE39E772C180E86039B2783A2EC07A28FB5C55DF06F4C52C9DE2BCBF6955817183995497CEA956AE515D2261898FA051015728E5A8AACAA68FFFFFFFFFFFFFFFF
g = 2
params = dh.DHParameterNumbers(p, g).parameters()

def setup(user):
    private_key = params.generate_private_key()
    public_key = private_key.public_key()

    skey = open(f"{user}.sk", "wb")
    pkey = open(f"{user}.pk", "wb")
    skey.write(
        private_key.private_bytes(
            encoding=Encoding.PEM,
            format=PrivateFormat.PKCS8,
            encryption_algorithm=NoEncryption(),
        )
    )
    pkey.write(
        public_key.public_bytes(
            encoding=Encoding.PEM, format=PublicFormat.SubjectPublicKeyInfo
        )
    )
    skey.close()
    pkey.close()


def encrypt(user, filename):
    peer_pkey_file = open(f"{user}.pk", "rb")
    if peer_pkey_file is None:
        print(f"No public key found for {user}")
        return

    file_r = open(filename, "rb")
    if file_r is None:
        print(f"File {filename} not found")
        return


    skey = params.generate_private_key()
    peer_pkey = serialization.load_pem_public_key(peer_pkey_file.read())
    if not isinstance(peer_pkey, dh.DHPublicKey):
        print("Error: Invalid key type")
        return

    shared_key = skey.exchange(peer_pkey)
    derived_key = HKDF(
        algorithm=hashes.SHA256(), length=32, salt=None, info=None
    ).derive(shared_key)

    message = file_r.read()
    chacha = ChaCha20Poly1305(derived_key)
    nonce = os.urandom(12)
    ciphertext = chacha.encrypt(nonce, message, None)
    file_w = open(f"{filename}.enc", "wb")
    nonce_ciphertext = nonce + ciphertext
    pkey_bytes = skey.public_key().public_bytes(
        Encoding.DER, PublicFormat.SubjectPublicKeyInfo
    )
    to_send = mkpair(pkey_bytes, nonce_ciphertext)
    file_w.write(to_send)

    file_r.close()
    file_w.close()
    peer_pkey_file.close()


def decrypt(user, filename):
    skey_file = open(f"{user}.sk", "rb")
    if skey_file is None:
        print(f"No private key found for {user}")
        return

    file_r = open(filename, "rb")
    if file_r is None:
        print(f"File {filename} not found")
        return

    message_bytes = file_r.read()
    peer_pkey_bytes, ciphertext = unpair(message_bytes)

    skey = serialization.load_pem_private_key(skey_file.read(), None)
    peer_pkey = serialization.load_der_public_key(peer_pkey_bytes)
    if not isinstance(skey, dh.DHPrivateKey) or not isinstance(
        peer_pkey, dh.DHPublicKey
    ):
        print("Error: Invalid key type")
        return

    shared_key = skey.exchange(peer_pkey)
    derived_key = HKDF(
        algorithm=hashes.SHA256(), length=32, salt=None, info=None
    ).derive(shared_key)

    nonce = ciphertext[:12]
    encrypted_message = ciphertext[12:]
    chacha = ChaCha20Poly1305(derived_key)
    message = chacha.decrypt(nonce, encrypted_message, None)
    file_w = open(f"{filename}.dec", "wb")
    file_w.write(message)
    file_r.close()
    file_w.close()
    skey_file.close()


def mkpair(x, y):
    """produz uma byte-string contendo o tuplo '(x,y)' ('x' e 'y' são byte-strings)"""
    len_x = len(x)
    len_x_bytes = len_x.to_bytes(2, "little")
    return len_x_bytes + x + y


def unpair(xy):
    """extrai componentes de um par codificado com 'mkpair'"""
    len_x = int.from_bytes(xy[:2], "little")
    x = xy[2 : len_x + 2]
    y = xy[len_x + 2 :]
    return x, y


def main():
    parser = argparse.ArgumentParser(
        description="Program to perform operations using ChaCha20 cipher on files",
    )

    subparsers = parser.add_subparsers(dest="operation", help="Operation to perform")

    # Setup subcommand
    setup_parser = subparsers.add_parser(
        "setup", help="Setup diffie helman key pair for user"
    )
    setup_parser.add_argument("user", help="User for which to setup the key pair")

    # Encrypt subcommand
    enc_parser = subparsers.add_parser("enc", help="Encrypt a file")
    enc_parser.add_argument("user", help="User for which to encrypt the file")
    enc_parser.add_argument("file", help="File to encrypt")

    # Decrypt subcommand
    decrypt_parser = subparsers.add_parser("dec", help="Decrypt a file")
    decrypt_parser.add_argument("user", help="Target user for the decryption")
    decrypt_parser.add_argument("file", help="File to be decrypted")

    args = parser.parse_args()
    match args.operation:
        case "setup":
            user = args.user
            setup(user)
        case "enc":
            user = args.user
            file = args.file
            encrypt(user, file)
        case "dec":
            user = args.user
            file = args.file
            decrypt(user, file)


if __name__ == "__main__":
    main()