#!/usr/bin/env python3
"""
node_kem.py
Montana Protocol — ML-KEM-768 Post-Quantum Key Exchange for Nodes

Постквантовый обмен ключами между узлами:
- ML-KEM-768 (FIPS 203) для инкапсуляции ключей
- AES-256-GCM для симметричного шифрования
- ML-DSA-65 для аутентификации

Защита от:
- Квантовых атак (Shor's algorithm, Grover's algorithm)
- Harvest now, decrypt later
- Man-in-the-middle (через ML-DSA-65 подписи)
"""

import hashlib
import os
import json
import asyncio
import logging
from pathlib import Path
from typing import Optional, Tuple, Dict, Any
from dataclasses import dataclass
from datetime import datetime, timezone

logger = logging.getLogger(__name__)

# ML-KEM-768 через kyber-py
try:
    from kyber_py.ml_kem import ML_KEM_768
    HAS_KYBER = True
except ImportError:
    HAS_KYBER = False
    logger.warning("kyber-py not installed: pip install kyber-py")

# AES-GCM
try:
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    HAS_CRYPTO = True
except ImportError:
    HAS_CRYPTO = False
    logger.warning("cryptography not installed: pip install cryptography")

# ML-DSA-65 для аутентификации
try:
    from dilithium_py.ml_dsa import ML_DSA_65
    HAS_DILITHIUM = True
except ImportError:
    HAS_DILITHIUM = False
    logger.warning("dilithium-py not installed: pip install dilithium-py")


# ═══════════════════════════════════════════════════════════════════════════════
#                              КОНСТАНТЫ
# ═══════════════════════════════════════════════════════════════════════════════

# ML-KEM-768 (FIPS 203)
PUBLIC_KEY_SIZE = 1184
SECRET_KEY_SIZE = 2400
CIPHERTEXT_SIZE = 1088
SHARED_SECRET_SIZE = 32
KEYPAIR_SEED_SIZE = 48

# Порт для постквантового обмена ключами
PQ_PORT = 19334

# Таймауты
HANDSHAKE_TIMEOUT = 30.0
MESSAGE_TIMEOUT = 60.0

# Версия протокола
PROTOCOL_VERSION = "PQ-1.0"


# ═══════════════════════════════════════════════════════════════════════════════
#                              ML-KEM-768 ФУНКЦИИ
# ═══════════════════════════════════════════════════════════════════════════════

def check_dependencies() -> bool:
    """Проверка зависимостей"""
    if not HAS_KYBER:
        return False
    if not HAS_CRYPTO:
        return False
    if not HAS_DILITHIUM:
        return False
    return True


def derive_kem_seed(private_key_dsa: bytes) -> bytes:
    """
    Деривирует seed для ML-KEM-768 из ML-DSA-65 приватного ключа узла.

    Это позволяет детерминированно генерировать KEM ключи
    для каждого узла на основе его identity.
    """
    # HKDF-подобная деривация
    salt = b"MONTANA_KEM_NODE_V1"

    # SHA3-256 хэш для деривации (не SHA-256, для квантовой стойкости)
    h = hashlib.sha3_256()
    h.update(salt)
    h.update(private_key_dsa)
    derived = h.digest()

    # Расширяем до 48 байт
    h2 = hashlib.sha3_256()
    h2.update(derived)
    h2.update(b"\x01")
    extended = derived + h2.digest()[:16]

    return extended[:KEYPAIR_SEED_SIZE]


def generate_kem_keypair(seed: bytes) -> Optional[Tuple[bytes, bytes]]:
    """
    Генерирует ML-KEM-768 ключевую пару детерминированно из seed.

    Returns:
        (secret_key, public_key) или None при ошибке
    """
    if not HAS_KYBER:
        logger.error("kyber-py not available")
        return None

    if len(seed) != KEYPAIR_SEED_SIZE:
        logger.error(f"Invalid seed size: {len(seed)}, expected {KEYPAIR_SEED_SIZE}")
        return None

    try:
        ML_KEM_768.set_drbg_seed(seed)
        public_key, secret_key = ML_KEM_768.keygen()
        return (secret_key, public_key)
    except Exception as e:
        logger.error(f"KEM keypair generation failed: {e}")
        return None


def encapsulate(public_key: bytes) -> Optional[Tuple[bytes, bytes]]:
    """
    Инкапсуляция для получения shared secret.

    Returns:
        (ciphertext, shared_secret) или None
    """
    if not HAS_KYBER:
        return None

    if len(public_key) != PUBLIC_KEY_SIZE:
        logger.error(f"Invalid public key size: {len(public_key)}")
        return None

    try:
        shared_secret, ciphertext = ML_KEM_768.encaps(public_key)
        return (ciphertext, shared_secret)
    except Exception as e:
        logger.error(f"Encapsulation failed: {e}")
        return None


def decapsulate(ciphertext: bytes, secret_key: bytes) -> Optional[bytes]:
    """
    Декапсуляция для восстановления shared secret.

    Returns:
        shared_secret (32 bytes) или None
    """
    if not HAS_KYBER:
        return None

    if len(ciphertext) != CIPHERTEXT_SIZE:
        logger.error(f"Invalid ciphertext size: {len(ciphertext)}")
        return None

    if len(secret_key) != SECRET_KEY_SIZE:
        logger.error(f"Invalid secret key size: {len(secret_key)}")
        return None

    try:
        shared_secret = ML_KEM_768.decaps(secret_key, ciphertext)
        return shared_secret
    except Exception as e:
        logger.error(f"Decapsulation failed: {e}")
        return None


# ═══════════════════════════════════════════════════════════════════════════════
#                              AES-256-GCM
# ═══════════════════════════════════════════════════════════════════════════════

def aes_encrypt(data: bytes, key: bytes) -> Optional[bytes]:
    """
    AES-256-GCM шифрование.

    Returns:
        nonce (12) + ciphertext + tag (16)
    """
    if not HAS_CRYPTO:
        return None

    if len(key) != 32:
        logger.error(f"Invalid key size: {len(key)}")
        return None

    try:
        nonce = os.urandom(12)
        aesgcm = AESGCM(key)
        ciphertext = aesgcm.encrypt(nonce, data, None)
        return nonce + ciphertext
    except Exception as e:
        logger.error(f"AES encryption failed: {e}")
        return None


def aes_decrypt(encrypted: bytes, key: bytes) -> Optional[bytes]:
    """
    AES-256-GCM расшифровка.

    Args:
        encrypted: nonce (12) + ciphertext + tag (16)
    """
    if not HAS_CRYPTO:
        return None

    if len(key) != 32:
        logger.error(f"Invalid key size: {len(key)}")
        return None

    if len(encrypted) < 28:  # 12 nonce + 16 tag минимум
        logger.error("Encrypted data too short")
        return None

    try:
        nonce = encrypted[:12]
        ciphertext = encrypted[12:]
        aesgcm = AESGCM(key)
        return aesgcm.decrypt(nonce, ciphertext, None)
    except Exception as e:
        logger.error(f"AES decryption failed: {e}")
        return None


# ═══════════════════════════════════════════════════════════════════════════════
#                              SECURE CHANNEL
# ═══════════════════════════════════════════════════════════════════════════════

@dataclass
class SecureChannel:
    """
    Защищённый канал между двумя узлами.

    Использует:
    - ML-KEM-768 для key exchange
    - AES-256-GCM для шифрования
    - ML-DSA-65 для аутентификации
    """

    local_address: str   # Наш адрес (mt...)
    peer_address: str    # Адрес партнёра (mt...)
    shared_secret: bytes  # 32 bytes from ML-KEM-768
    established_at: datetime

    # Счётчики для nonce (предотвращение replay)
    send_counter: int = 0
    recv_counter: int = 0

    def derive_key(self, sender_address: str, counter: int) -> bytes:
        """
        Деривирует ключ на основе отправителя.

        Ключ детерминирован: оба узла получат одинаковый ключ
        для сообщения от конкретного отправителя.
        """
        h = hashlib.sha3_256()
        h.update(self.shared_secret)
        h.update(sender_address.encode())
        h.update(counter.to_bytes(8, 'big'))
        return h.digest()

    def encrypt(self, data: bytes) -> Optional[bytes]:
        """Шифрует данные для отправки"""
        self.send_counter += 1
        key = self.derive_key(self.local_address, self.send_counter)

        encrypted = aes_encrypt(data, key)
        if encrypted:
            # Добавляем counter в начало
            return self.send_counter.to_bytes(8, 'big') + encrypted
        return None

    def decrypt(self, data: bytes) -> Optional[bytes]:
        """Расшифровывает полученные данные"""
        if len(data) < 8:
            return None

        counter = int.from_bytes(data[:8], 'big')

        # Защита от replay
        if counter <= self.recv_counter:
            logger.warning(f"Replay attack detected: {counter} <= {self.recv_counter}")
            return None

        # Ключ деривируется от адреса отправителя (peer)
        key = self.derive_key(self.peer_address, counter)
        decrypted = aes_decrypt(data[8:], key)

        if decrypted:
            self.recv_counter = counter

        return decrypted


# ═══════════════════════════════════════════════════════════════════════════════
#                              NODE KEM MANAGER
# ═══════════════════════════════════════════════════════════════════════════════

class NodeKEMManager:
    """
    Менеджер ML-KEM-768 для узла Montana.

    Управляет:
    - Генерацией KEM ключей из DSA identity
    - Установлением защищённых каналов
    - Хранением активных сессий
    """

    def __init__(self, node_address: str, dsa_private_key: bytes, dsa_public_key: bytes):
        """
        Args:
            node_address: Адрес узла (mt...)
            dsa_private_key: ML-DSA-65 приватный ключ (4032 bytes)
            dsa_public_key: ML-DSA-65 публичный ключ (1952 bytes)
        """
        self.node_address = node_address
        self.dsa_private_key = dsa_private_key
        self.dsa_public_key = dsa_public_key

        # Деривируем KEM ключи из DSA identity
        seed = derive_kem_seed(dsa_private_key)
        keys = generate_kem_keypair(seed)

        if keys:
            self.kem_secret_key, self.kem_public_key = keys
            logger.info(f"[NodeKEM] Initialized for {node_address[:16]}...")
        else:
            self.kem_secret_key = None
            self.kem_public_key = None
            logger.error("[NodeKEM] Failed to generate KEM keys")

        # Активные защищённые каналы
        self.channels: Dict[str, SecureChannel] = {}

    def get_handshake_request(self) -> Dict[str, Any]:
        """
        Создаёт запрос на установление защищённого канала.

        Returns:
            {
                "type": "kem_handshake",
                "version": "PQ-1.0",
                "node_address": "mt...",
                "kem_public_key": "hex...",  # 1184 bytes
                "dsa_public_key": "hex...",  # 1952 bytes
                "timestamp": "iso...",
                "signature": "hex..."        # ML-DSA-65 подпись
            }
        """
        if not self.kem_public_key:
            return {"error": "KEM not initialized"}

        timestamp = datetime.now(timezone.utc).isoformat()

        message = f"KEM_HANDSHAKE:{self.node_address}:{timestamp}"
        signature = ML_DSA_65.sign(self.dsa_private_key, message.encode())

        return {
            "type": "kem_handshake",
            "version": PROTOCOL_VERSION,
            "node_address": self.node_address,
            "kem_public_key": self.kem_public_key.hex(),
            "dsa_public_key": self.dsa_public_key.hex(),
            "timestamp": timestamp,
            "signature": signature.hex()
        }

    def process_handshake_request(
        self,
        request: Dict[str, Any]
    ) -> Optional[Tuple[Dict[str, Any], SecureChannel]]:
        """
        Обрабатывает запрос handshake от другого узла.

        Returns:
            (response, channel) или None при ошибке
        """
        try:
            # Извлекаем данные
            peer_address = request["node_address"]
            peer_kem_pk = bytes.fromhex(request["kem_public_key"])
            peer_dsa_pk = bytes.fromhex(request["dsa_public_key"])
            timestamp = request["timestamp"]
            signature = bytes.fromhex(request["signature"])

            # Проверяем подпись (аутентификация)
            message = f"KEM_HANDSHAKE:{peer_address}:{timestamp}"
            if not ML_DSA_65.verify(peer_dsa_pk, message.encode(), signature):
                logger.warning(f"[NodeKEM] Invalid signature from {peer_address[:16]}...")
                return None

            # Проверяем что адрес соответствует публичному ключу
            expected_address = "mt" + hashlib.sha256(peer_dsa_pk).hexdigest()[:40]
            if peer_address != expected_address:
                logger.warning(f"[NodeKEM] Address mismatch: {peer_address} != {expected_address}")
                return None

            # Инкапсуляция: создаём shared secret
            result = encapsulate(peer_kem_pk)
            if not result:
                return None

            ciphertext, shared_secret = result

            # Создаём канал
            channel = SecureChannel(
                local_address=self.node_address,
                peer_address=peer_address,
                shared_secret=shared_secret,
                established_at=datetime.now(timezone.utc)
            )

            # Сохраняем
            self.channels[peer_address] = channel

            # Подписываем ответ
            response_ts = datetime.now(timezone.utc).isoformat()
            response_msg = f"KEM_RESPONSE:{self.node_address}:{response_ts}"
            response_sig = ML_DSA_65.sign(self.dsa_private_key, response_msg.encode())

            response = {
                "type": "kem_response",
                "version": PROTOCOL_VERSION,
                "node_address": self.node_address,
                "ciphertext": ciphertext.hex(),  # 1088 bytes
                "dsa_public_key": self.dsa_public_key.hex(),
                "timestamp": response_ts,
                "signature": response_sig.hex()
            }

            logger.info(f"[NodeKEM] Channel established with {peer_address[:16]}...")
            return (response, channel)

        except Exception as e:
            logger.error(f"[NodeKEM] Handshake processing failed: {e}")
            return None

    def process_handshake_response(
        self,
        response: Dict[str, Any]
    ) -> Optional[SecureChannel]:
        """
        Обрабатывает ответ на наш handshake запрос.

        Returns:
            SecureChannel или None
        """
        try:
            # Извлекаем данные
            peer_address = response["node_address"]
            ciphertext = bytes.fromhex(response["ciphertext"])
            peer_dsa_pk = bytes.fromhex(response["dsa_public_key"])
            timestamp = response["timestamp"]
            signature = bytes.fromhex(response["signature"])

            # Проверяем подпись
            message = f"KEM_RESPONSE:{peer_address}:{timestamp}"
            if not ML_DSA_65.verify(peer_dsa_pk, message.encode(), signature):
                logger.warning(f"[NodeKEM] Invalid response signature from {peer_address[:16]}...")
                return None

            # Декапсуляция: восстанавливаем shared secret
            shared_secret = decapsulate(ciphertext, self.kem_secret_key)
            if not shared_secret:
                return None

            # Создаём канал
            channel = SecureChannel(
                local_address=self.node_address,
                peer_address=peer_address,
                shared_secret=shared_secret,
                established_at=datetime.now(timezone.utc)
            )

            # Сохраняем
            self.channels[peer_address] = channel

            logger.info(f"[NodeKEM] Channel established with {peer_address[:16]}...")
            return channel

        except Exception as e:
            logger.error(f"[NodeKEM] Response processing failed: {e}")
            return None

    def get_channel(self, peer_address: str) -> Optional[SecureChannel]:
        """Получить защищённый канал к узлу"""
        return self.channels.get(peer_address)

    def close_channel(self, peer_address: str):
        """Закрыть канал"""
        if peer_address in self.channels:
            del self.channels[peer_address]
            logger.info(f"[NodeKEM] Channel closed: {peer_address[:16]}...")


# ═══════════════════════════════════════════════════════════════════════════════
#                              ASYNC CLIENT/SERVER
# ═══════════════════════════════════════════════════════════════════════════════

class PQSecureClient:
    """
    Асинхронный клиент для постквантового обмена ключами.
    """

    def __init__(self, kem_manager: NodeKEMManager):
        self.kem = kem_manager

    async def establish_channel(
        self,
        host: str,
        port: int = PQ_PORT,
        timeout: float = HANDSHAKE_TIMEOUT
    ) -> Optional[SecureChannel]:
        """
        Устанавливает защищённый канал с узлом.
        """
        try:
            reader, writer = await asyncio.wait_for(
                asyncio.open_connection(host, port),
                timeout=timeout
            )

            # Отправляем handshake request
            request = self.kem.get_handshake_request()
            data = json.dumps(request).encode()
            writer.write(len(data).to_bytes(4, 'big'))
            writer.write(data)
            await writer.drain()

            # Получаем response
            length_bytes = await asyncio.wait_for(
                reader.read(4),
                timeout=timeout
            )
            length = int.from_bytes(length_bytes, 'big')
            response_data = await asyncio.wait_for(
                reader.read(length),
                timeout=timeout
            )

            response = json.loads(response_data.decode())

            if response.get("type") == "kem_response":
                channel = self.kem.process_handshake_response(response)
                writer.close()
                await writer.wait_closed()
                return channel

            writer.close()
            await writer.wait_closed()
            return None

        except Exception as e:
            logger.error(f"[PQClient] Failed to establish channel with {host}: {e}")
            return None


class PQSecureServer:
    """
    Асинхронный сервер для постквантового обмена ключами.
    """

    def __init__(self, kem_manager: NodeKEMManager):
        self.kem = kem_manager
        self._server = None

    async def _handle_client(
        self,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter
    ):
        """Обрабатывает handshake от клиента"""
        peer = writer.get_extra_info('peername')

        try:
            # Читаем request
            length_bytes = await asyncio.wait_for(
                reader.read(4),
                timeout=HANDSHAKE_TIMEOUT
            )
            length = int.from_bytes(length_bytes, 'big')
            request_data = await asyncio.wait_for(
                reader.read(length),
                timeout=HANDSHAKE_TIMEOUT
            )

            request = json.loads(request_data.decode())

            if request.get("type") == "kem_handshake":
                result = self.kem.process_handshake_request(request)

                if result:
                    response, channel = result
                    data = json.dumps(response).encode()
                    writer.write(len(data).to_bytes(4, 'big'))
                    writer.write(data)
                    await writer.drain()
                    logger.info(f"[PQServer] Handshake complete with {peer}")
                else:
                    error = {"type": "error", "message": "Handshake failed"}
                    data = json.dumps(error).encode()
                    writer.write(len(data).to_bytes(4, 'big'))
                    writer.write(data)
                    await writer.drain()

        except Exception as e:
            logger.error(f"[PQServer] Error handling {peer}: {e}")
        finally:
            writer.close()
            await writer.wait_closed()

    async def start(self, host: str = "0.0.0.0", port: int = PQ_PORT):
        """Запускает сервер"""
        self._server = await asyncio.start_server(
            self._handle_client,
            host, port
        )
        logger.info(f"[PQServer] Started on {host}:{port}")

        async with self._server:
            await self._server.serve_forever()

    def stop(self):
        """Останавливает сервер"""
        if self._server:
            self._server.close()


# ═══════════════════════════════════════════════════════════════════════════════
#                              ТЕСТИРОВАНИЕ
# ═══════════════════════════════════════════════════════════════════════════════

def test_node_kem():
    """Тест ML-KEM-768 для узлов"""
    print("=" * 60)
    print("Montana Protocol — ML-KEM-768 Node Test")
    print("=" * 60 + "\n")

    if not check_dependencies():
        print("Missing dependencies. Install:")
        print("  pip install kyber-py cryptography dilithium-py")
        return False

    # Создаём два узла
    print("Generating Node 1 (Amsterdam)...")
    pk1, sk1 = ML_DSA_65.keygen()
    addr1 = "mt" + hashlib.sha256(pk1).hexdigest()[:40]

    print("Generating Node 2 (Moscow)...")
    pk2, sk2 = ML_DSA_65.keygen()
    addr2 = "mt" + hashlib.sha256(pk2).hexdigest()[:40]

    print(f"\nNode 1: {addr1[:20]}...")
    print(f"Node 2: {addr2[:20]}...")

    # Инициализируем KEM менеджеры
    print("\nInitializing KEM managers...")
    kem1 = NodeKEMManager(addr1, sk1, pk1)
    kem2 = NodeKEMManager(addr2, sk2, pk2)

    if not kem1.kem_public_key or not kem2.kem_public_key:
        print("Failed to initialize KEM")
        return False

    # Handshake
    print("\nPerforming handshake...")

    # Node 1 создаёт запрос
    request = kem1.get_handshake_request()
    print(f"  Request size: {len(json.dumps(request))} bytes")

    # Node 2 обрабатывает и отвечает
    result = kem2.process_handshake_request(request)
    if not result:
        print("Handshake failed at Node 2")
        return False

    response, channel2 = result
    print(f"  Response size: {len(json.dumps(response))} bytes")

    # Node 1 обрабатывает ответ
    channel1 = kem1.process_handshake_response(response)
    if not channel1:
        print("Handshake failed at Node 1")
        return False

    print("\nHandshake complete!")
    print(f"  Shared secret match: {channel1.shared_secret == channel2.shared_secret}")

    # Тест шифрования
    print("\nTesting encrypted communication...")

    test_message = b"Hello from Node 1! This is a test message."

    # Node 1 шифрует
    encrypted = channel1.encrypt(test_message)
    print(f"  Original: {len(test_message)} bytes")
    print(f"  Encrypted: {len(encrypted)} bytes")

    # Node 2 расшифровывает
    decrypted = channel2.decrypt(encrypted)
    print(f"  Decrypted: {decrypted}")

    if decrypted == test_message:
        print("\nSUCCESS: Message decrypted correctly")
    else:
        print("\nFAIL: Message mismatch")
        return False

    # Тест replay protection
    print("\nTesting replay protection...")
    replay_result = channel2.decrypt(encrypted)
    if replay_result is None:
        print("  Replay attack blocked correctly")
    else:
        print("  FAIL: Replay attack not detected!")
        return False

    print("\n" + "=" * 60)
    print("ALL TESTS PASSED")
    print("=" * 60)
    return True


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    test_node_kem()