# compatibile Windows 11
# compatibile Ubuntu 24.10
# compatibile python 3.12.7

import os
from typing import Tuple
from cryptography.hazmat.primitives.asymmetric import dh
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization


def generate_parameters(key_size: int = 2048) -> dh.DHParameters:
    """
    Genera i parametri Diffie-Hellman.

    Args:
        key_size: La dimensione della chiave in bit (default: 2048).

    Returns:
        I parametri Diffie-Hellman.
    """
    parameters = dh.generate_parameters(generator=2, key_size=key_size,
                                        backend=default_backend())
    return parameters

def generate_keypair(parameters: dh.DHParameters) -> Tuple[dh.DHPrivateKey, dh.DHPublicKey]:
    """
    Genera una coppia di chiavi Diffie-Hellman (privata, pubblica).

    Args:
        parameters: I parametri Diffie-Hellman.

    Returns:
        Una tupla (private_key, public_key).
    """
    private_key = parameters.generate_private_key()
    public_key = private_key.public_key()
    return private_key, public_key

def compute_shared_secret(private_key: dh.DHPrivateKey, public_key: dh.DHPublicKey) -> bytes:
    """
    Calcola il segreto condiviso usando la chiave privata e la chiave pubblica dell'altra parte.

    Args:
        private_key: La propria chiave privata.
        public_key: La chiave pubblica dell'altra parte.

    Returns:
        Il segreto condiviso (bytes).
    """
    shared_secret = private_key.exchange(public_key)
    return shared_secret

def derive_key(shared_secret: bytes) -> bytes:
    """
    Deriva una chiave simmetrica dal segreto condiviso usando HKDF.

    Args:
        shared_secret: Il segreto condiviso.

    Returns:
        La chiave simmetrica derivata (bytes).
    """
    # Usa HKDF per derivare una chiave simmetrica dal segreto condiviso
    hkdf = HKDF(
        algorithm=hashes.SHA256(),
        length=32,  # Lunghezza della chiave desiderata (es. 32 byte per AES-256)
        salt=None,  # Un salt casuale aumenterebbe la sicurezza
        info=b'handshake data',  # Contesto dell'applicazione (può essere qualsiasi cosa)
        backend=default_backend()
    )
    key = hkdf.derive(shared_secret)
    return key

def serialize_public_key(public_key: dh.DHPublicKey) -> bytes:
    """Serializza la chiave pubblica in formato PEM."""
    pem = public_key.public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo
    )
    return pem

def deserialize_public_key(pem: bytes) -> dh.DHPublicKey:
    """Deserializza la chiave pubblica da formato PEM."""
    public_key = serialization.load_pem_public_key(
        pem,
        backend=default_backend()
    )
    return public_key

def serialize_private_key(private_key: dh.DHPrivateKey, password: str = None) -> bytes:
    """Serializza la chiave privata in formato PEM, con opzione per la password."""
    if password:
        encryption = serialization.BestAvailableEncryption(password.encode())
    else:
        encryption = serialization.NoEncryption()
    pem = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=encryption
    )
    return pem

def deserialize_private_key(pem: bytes, password: str = None) -> dh.DHPrivateKey:
    """Deserializza la chiave privata da formato PEM, con password opzionale."""
    private_key = serialization.load_pem_private_key(
        pem,
        password=password.encode() if password else None,
        backend=default_backend()
    )
    return private_key


def clear_screen():
    """Pulisce lo schermo della console (sia Windows che Unix-like)."""
    os.system('cls' if os.name == 'nt' else 'clear')

def main():
    """Funzione principale del programma dimostrativo."""

    clear_screen()
    print("=" * 60)
    print("   Dimostrazione dello scambio di chiavi Diffie-Hellman con 'cryptography'")
    print("=" * 60)

    # 1. Generazione dei parametri (comuni)
    print("\n1. Generazione dei parametri...")
    key_size = 0
    while key_size not in [1024, 2048, 3072]:
        try:
            key_size = int(input("   Inserisci la dimensione della chiave (1024, 2048, 3072): "))
            if key_size not in [1024, 2048, 3072]:
                print("   Errore: la dimensione deve essere 1024, 2048 o 3072.")
        except ValueError:
            print("   Errore: inserisci un numero intero.")

    parameters = generate_parameters(key_size)
    print("   Parametri generati.")

    # 2. Generazione delle chiavi di Alice
    print("\n2. Generazione delle chiavi di Alice...")
    alice_private_key, alice_public_key = generate_keypair(parameters)
    print("   Chiavi di Alice generate.")

    # 3. Generazione delle chiavi di Bob
    print("\n3. Generazione delle chiavi di Bob...")
    bob_private_key, bob_public_key = generate_keypair(parameters)
    print("   Chiavi di Bob generate.")

    # 4. Scambio delle chiavi pubbliche (simulato)
    print("\n4. Scambio delle chiavi pubbliche (simulato)...")
    print("   Alice invia la sua chiave pubblica a Bob.")
    print("   Bob invia la sua chiave pubblica ad Alice.")

    # Serializzazione delle chiavi pubbliche per lo scambio
    alice_public_pem = serialize_public_key(alice_public_key)
    bob_public_pem = serialize_public_key(bob_public_key)

    # Deserializzazione delle chiavi pubbliche ricevute
    received_alice_public_key = deserialize_public_key(alice_public_pem)
    received_bob_public_key = deserialize_public_key(bob_public_pem)
    print("\n Chiavi pubbliche serializzate e deserializzate per lo scambio")


    # 5. Calcolo del segreto condiviso
    print("\n5. Calcolo del segreto condiviso...")
    alice_shared_secret = compute_shared_secret(alice_private_key, received_bob_public_key)
    bob_shared_secret = compute_shared_secret(bob_private_key, received_alice_public_key)
    print(f"   Segreto condiviso da Alice (bytes): {alice_shared_secret.hex()}")
    print(f"   Segreto condiviso da Bob (bytes): {bob_shared_secret.hex()}")


    # 6. Derivazione di una chiave simmetrica
    print("\n6. Derivazione di una chiave simmetrica dal segreto condiviso...")
    alice_key = derive_key(alice_shared_secret)
    bob_key = derive_key(bob_shared_secret)
    print(f"   Chiave simmetrica derivata da Alice: {alice_key.hex()}")
    print(f"   Chiave simmetrica derivata da Bob: {bob_key.hex()}")

    # 7. Verifica
    if alice_key == bob_key:
        print("\n   [SUCCESS] Lo scambio di chiavi ha avuto successo!")
        print("   Alice e Bob hanno derivato la stessa chiave simmetrica.")
    else:
        print("\n   [ERRORE] Lo scambio di chiavi è fallito!")
        print("   Alice e Bob NON hanno derivato la stessa chiave simmetrica.")

        # Serializzazione delle chiavi private per debugging/dimostrazione (opzionale, in un'applicazione reale le chiavi private non verrebbero mai stampate)
    alice_private_pem = serialize_private_key(alice_private_key, password="password")
    bob_private_pem = serialize_private_key(bob_private_key, password="password")
    print("\n Chiave privata di Alice (PEM, protetta):")
    print(alice_private_pem.decode())
    print("\n Chiave privata di Bob (PEM, protetta):")
    print(bob_private_pem.decode())


    print("=" * 60)

if __name__ == "__main__":
    main()