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

# client.py (versione con scambio di chiavi)
import socket
import struct
import os
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding

# --- Configurazione (deve corrispondere al server) ---
MULTICAST_GROUP = '224.1.1.1'
MULTICAST_PORT = 5007
FRAGMENT_SIZE = 256
MAX_PACKET_SIZE = 65507

# --- Funzioni di utilità (dovrebbero essere in utils.py, importate qui) ---
# (Assicurati che utils.py contenga create_tlv, encrypt_data, fragment_data)

def create_tlv(type, value):
    """Crea un messaggio TLV."""
    type_bytes = struct.pack('!H', type)  # Tipo: unsigned short (2 byte, big-endian)
    length_bytes = struct.pack('!H', len(value))  # Lunghezza: unsigned short
    return type_bytes + length_bytes + value

def parse_tlv(data):
    """Analizza un messaggio TLV. Restituisce (type, value, resto_del_buffer)"""
    if len(data) < 4:  # Almeno tipo (2) + lunghezza (2)
        raise ValueError("Dati TLV insufficienti.")
    type = struct.unpack('!H', data[:2])[0]
    length = struct.unpack('!H', data[2:4])[0]
    if len(data) < 4 + length:
         raise ValueError("Lunghezza TLV non corrisponde ai dati.")
    value = data[4:4 + length]
    remaining_data = data[4 + length:]
    return type, value, remaining_data


# --- Funzioni di crittografia ---

def encrypt_data(data, public_key):
    """Crittografa i dati con la chiave pubblica RSA (OAEP)."""
    ciphertext = public_key.encrypt(
        data,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None
        )
    )
    return ciphertext

def decrypt_data(ciphertext, private_key):
     """Decrittografa i dati con la chiave privata RSA (OAEP)."""
     try:
        plaintext = private_key.decrypt(
            ciphertext,
            padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None
            )
        )
        return plaintext
     except Exception as e:
         print(f"Errore durante la decrittazione: {e}")
         return None


# --- Frammentazione e deframmentazione ---

def fragment_data(data, fragment_size):
    """Divide i dati in frammenti di dimensione specificata."""
    return [data[i:i+fragment_size] for i in range(0, len(data), fragment_size)]

def defragment_data(fragments):
    """Ricompone i dati dai frammenti (senza controllo d'ordine!)."""
    return b''.join(fragments)


def request_public_key(sock):
    """Richiede la chiave pubblica al server e la restituisce."""
    tlv_request = create_tlv(4, b"")  # Tipo 4: Richiesta chiave pubblica (payload vuoto)
    sock.sendto(tlv_request, (MULTICAST_GROUP, MULTICAST_PORT))
    print("Richiesta chiave pubblica inviata...")

    # Attendi la risposta (con timeout)
    sock.settimeout(5.0)  # Timeout di 5 secondi per la ricezione della chiave
    try:
        data, address = sock.recvfrom(MAX_PACKET_SIZE)
        type, value, _ = parse_tlv(data)
        if type == 5:  # Tipo 5: Risposta chiave pubblica
            public_key = serialization.load_pem_public_key(
                value,
                backend=default_backend()
            )
            print("Chiave pubblica ricevuta.")
            return public_key
        else:
            print("Risposta non valida dal server.")
            return None
    except socket.timeout:
        print("Timeout durante la ricezione della chiave pubblica.")
        return None
    except Exception as e:
        print("Errore", e)
        return None



def send_message(message, public_key, sock):
    """Invia un messaggio, gestisce frammentazione pre/post crittografia, TLV."""
    message_bytes = message.encode('utf-8')
    key_size = public_key.key_size // 8
    max_block_size = key_size - 2 * hashes.SHA256.digest_size - 2

    pre_fragments = [message_bytes[i:i + max_block_size] for i in range(0, len(message_bytes), max_block_size)] \
                    if len(message_bytes) > max_block_size else [message_bytes]

    encrypted_tlvs = [create_tlv(3, encrypt_data(pre_frag, public_key)) for pre_frag in pre_fragments]

    msg_id = os.urandom(4)
    combined_data = b"".join(encrypted_tlvs)

    if len(combined_data) > FRAGMENT_SIZE:
        udp_fragments = fragment_data(combined_data, FRAGMENT_SIZE)
        total_udp_fragments = len(udp_fragments)
        for seq_num, udp_frag in enumerate(udp_fragments):
            tlv_data = struct.pack('!I', int.from_bytes(msg_id, byteorder='big')) + struct.pack('!H', seq_num) + struct.pack('!H', total_udp_fragments) + udp_frag
            tlv_udp = create_tlv(1, tlv_data)
            sock.sendto(tlv_udp, (MULTICAST_GROUP, MULTICAST_PORT))
            print(f"Inviato frammento UDP {seq_num + 1}/{total_udp_fragments}")
    else:
        tlv_data = struct.pack('!I', int.from_bytes(msg_id, byteorder='big')) + struct.pack('!H', 0) + struct.pack('!H', 1) + combined_data
        tlv_udp = create_tlv(1, tlv_data)
        sock.sendto(tlv_udp, (MULTICAST_GROUP, MULTICAST_PORT))
        print("Inviato messaggio (senza frammentazione UDP)")
    print("Messaggio inviato con successo")

def run_client():
    """Funzione principale del client."""
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    ttl = struct.pack('b', 1)
    sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, ttl)

    # Richiedi la chiave pubblica
    public_key = request_public_key(sock)
    if public_key is None:
        sock.close()
        return  # Esci se non si ottiene la chiave

    while True:
        message = input("Inserisci il messaggio da inviare (o 'exit' per uscire): ")
        if message.lower() == 'exit':
            break
        send_message(message, public_key, sock)

    sock.close()

if __name__ == "__main__":
    run_client()