862 lines
33 KiB
Python
862 lines
33 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
node_tls.py
|
||
Montana Protocol — Hybrid TLS + ML-KEM-768 Шифрование между узлами
|
||
|
||
Обеспечивает защищённую связь между узлами Montana:
|
||
- TLS 1.3 шифрование всего трафика (классическая защита)
|
||
- ML-KEM-768 post-quantum key exchange (FIPS 203)
|
||
- Взаимная аутентификация узлов через ML-DSA-65
|
||
|
||
HYBRID ЗАЩИТА:
|
||
- TLS 1.3: защита от классических атак
|
||
- ML-KEM-768: защита от квантовых атак
|
||
- Данные защищены даже если одна из систем скомпрометирована
|
||
|
||
Защита от:
|
||
- Квантовых атак (Shor's algorithm)
|
||
- Harvest now, decrypt later
|
||
- Man-in-the-middle атак
|
||
- Перехвата трафика
|
||
"""
|
||
|
||
import ssl
|
||
import socket
|
||
import asyncio
|
||
import logging
|
||
import hashlib
|
||
import json
|
||
from pathlib import Path
|
||
from typing import Optional, Dict, Any, Tuple
|
||
from datetime import datetime, timezone
|
||
from dataclasses import dataclass
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# ML-KEM-768 Post-Quantum Key Exchange
|
||
try:
|
||
from node_kem import (
|
||
NodeKEMManager,
|
||
SecureChannel as PQSecureChannel,
|
||
check_dependencies as check_pq_dependencies,
|
||
PQ_PORT
|
||
)
|
||
HAS_PQ = check_pq_dependencies()
|
||
except ImportError:
|
||
HAS_PQ = False
|
||
logger.warning("node_kem not available, post-quantum disabled")
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════════════
|
||
# КОНФИГУРАЦИЯ
|
||
# ═══════════════════════════════════════════════════════════════════════════════
|
||
|
||
@dataclass
|
||
class TLSConfig:
|
||
"""Конфигурация TLS"""
|
||
|
||
# Порт для защищённой связи
|
||
SECURE_PORT: int = 19333
|
||
|
||
# Таймауты
|
||
CONNECT_TIMEOUT: float = 10.0
|
||
READ_TIMEOUT: float = 30.0
|
||
|
||
# SSL/TLS параметры
|
||
TLS_VERSION = ssl.TLSVersion.TLSv1_3
|
||
VERIFY_MODE = ssl.CERT_OPTIONAL # CERT_REQUIRED для production
|
||
|
||
# Пути к сертификатам (генерируются автоматически)
|
||
CERTS_DIR: str = "certs"
|
||
CERT_FILE: str = "node.crt"
|
||
KEY_FILE: str = "node.key"
|
||
CA_FILE: str = "montana_ca.crt"
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════════════
|
||
# ГЕНЕРАЦИЯ СЕРТИФИКАТОВ
|
||
# ═══════════════════════════════════════════════════════════════════════════════
|
||
|
||
class CertificateManager:
|
||
"""
|
||
Менеджер сертификатов для узлов Montana
|
||
|
||
Генерирует self-signed сертификаты с привязкой к
|
||
криптографическому адресу узла (mt...).
|
||
"""
|
||
|
||
def __init__(self, data_dir: Path = None):
|
||
self.data_dir = data_dir or Path(__file__).parent
|
||
self.certs_dir = self.data_dir / TLSConfig.CERTS_DIR
|
||
self.certs_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
def generate_self_signed_cert(
|
||
self,
|
||
node_address: str,
|
||
node_name: str,
|
||
days_valid: int = 365
|
||
) -> Tuple[Path, Path]:
|
||
"""
|
||
Генерирует self-signed сертификат для узла
|
||
|
||
Args:
|
||
node_address: Криптографический адрес узла (mt...)
|
||
node_name: Имя узла (amsterdam, moscow, etc.)
|
||
days_valid: Срок действия в днях
|
||
|
||
Returns:
|
||
(cert_path, key_path)
|
||
"""
|
||
try:
|
||
from cryptography import x509
|
||
from cryptography.x509.oid import NameOID
|
||
from cryptography.hazmat.primitives import hashes
|
||
from cryptography.hazmat.primitives.asymmetric import rsa
|
||
from cryptography.hazmat.primitives import serialization
|
||
from datetime import timedelta
|
||
|
||
# Генерируем RSA ключ для TLS (ML-DSA-65 не поддерживается в TLS пока)
|
||
key = rsa.generate_private_key(
|
||
public_exponent=65537,
|
||
key_size=4096,
|
||
)
|
||
|
||
# Создаём сертификат
|
||
subject = issuer = x509.Name([
|
||
x509.NameAttribute(NameOID.COUNTRY_NAME, "NL"),
|
||
x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "Montana"),
|
||
x509.NameAttribute(NameOID.LOCALITY_NAME, node_name),
|
||
x509.NameAttribute(NameOID.ORGANIZATION_NAME, "Montana Protocol"),
|
||
x509.NameAttribute(NameOID.COMMON_NAME, f"{node_name}.efir.org"),
|
||
x509.NameAttribute(NameOID.USER_ID, node_address),
|
||
])
|
||
|
||
cert = x509.CertificateBuilder().subject_name(
|
||
subject
|
||
).issuer_name(
|
||
issuer
|
||
).public_key(
|
||
key.public_key()
|
||
).serial_number(
|
||
x509.random_serial_number()
|
||
).not_valid_before(
|
||
datetime.now(timezone.utc)
|
||
).not_valid_after(
|
||
datetime.now(timezone.utc) + timedelta(days=days_valid)
|
||
).add_extension(
|
||
x509.SubjectAlternativeName([
|
||
x509.DNSName(f"{node_name}.efir.org"),
|
||
x509.DNSName("localhost"),
|
||
]),
|
||
critical=False,
|
||
).sign(key, hashes.SHA256())
|
||
|
||
# Сохраняем
|
||
cert_path = self.certs_dir / f"{node_name}.crt"
|
||
key_path = self.certs_dir / f"{node_name}.key"
|
||
|
||
with open(cert_path, "wb") as f:
|
||
f.write(cert.public_bytes(serialization.Encoding.PEM))
|
||
|
||
with open(key_path, "wb") as f:
|
||
f.write(key.private_bytes(
|
||
encoding=serialization.Encoding.PEM,
|
||
format=serialization.PrivateFormat.TraditionalOpenSSL,
|
||
encryption_algorithm=serialization.NoEncryption()
|
||
))
|
||
|
||
logger.info(f"🔐 Сертификат создан: {cert_path}")
|
||
return cert_path, key_path
|
||
|
||
except ImportError:
|
||
logger.warning("⚠️ cryptography не установлена, используем openssl")
|
||
return self._generate_with_openssl(node_address, node_name, days_valid)
|
||
|
||
def _generate_with_openssl(
|
||
self,
|
||
node_address: str,
|
||
node_name: str,
|
||
days_valid: int
|
||
) -> Tuple[Path, Path]:
|
||
"""Генерация через openssl CLI"""
|
||
import subprocess
|
||
|
||
cert_path = self.certs_dir / f"{node_name}.crt"
|
||
key_path = self.certs_dir / f"{node_name}.key"
|
||
|
||
# Генерируем ключ и сертификат
|
||
cmd = [
|
||
"openssl", "req", "-x509", "-newkey", "rsa:4096",
|
||
"-keyout", str(key_path),
|
||
"-out", str(cert_path),
|
||
"-days", str(days_valid),
|
||
"-nodes",
|
||
"-subj", f"/CN={node_name}.efir.org/O=Montana Protocol/UID={node_address}"
|
||
]
|
||
|
||
result = subprocess.run(cmd, capture_output=True, text=True)
|
||
if result.returncode != 0:
|
||
raise RuntimeError(f"openssl error: {result.stderr}")
|
||
|
||
logger.info(f"🔐 Сертификат создан (openssl): {cert_path}")
|
||
return cert_path, key_path
|
||
|
||
def get_cert_paths(self, node_name: str) -> Tuple[Optional[Path], Optional[Path]]:
|
||
"""Получить пути к сертификатам узла"""
|
||
cert_path = self.certs_dir / f"{node_name}.crt"
|
||
key_path = self.certs_dir / f"{node_name}.key"
|
||
|
||
if cert_path.exists() and key_path.exists():
|
||
return cert_path, key_path
|
||
return None, None
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════════════
|
||
# TLS КОНТЕКСТ
|
||
# ═══════════════════════════════════════════════════════════════════════════════
|
||
|
||
class SecureContext:
|
||
"""
|
||
Создаёт SSL контекст для защищённой связи
|
||
"""
|
||
|
||
@staticmethod
|
||
def create_server_context(
|
||
cert_path: Path,
|
||
key_path: Path,
|
||
ca_path: Path = None
|
||
) -> ssl.SSLContext:
|
||
"""
|
||
Создаёт SSL контекст для сервера
|
||
|
||
Args:
|
||
cert_path: Путь к сертификату
|
||
key_path: Путь к приватному ключу
|
||
ca_path: Путь к CA сертификату (для mTLS)
|
||
"""
|
||
context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
|
||
context.minimum_version = TLSConfig.TLS_VERSION
|
||
|
||
# Загружаем сертификат
|
||
context.load_cert_chain(str(cert_path), str(key_path))
|
||
|
||
# CA для проверки клиентов (mTLS)
|
||
if ca_path and ca_path.exists():
|
||
context.load_verify_locations(str(ca_path))
|
||
context.verify_mode = ssl.CERT_REQUIRED
|
||
else:
|
||
context.verify_mode = ssl.CERT_OPTIONAL
|
||
|
||
# Безопасные настройки
|
||
context.set_ciphers('ECDHE+AESGCM:DHE+AESGCM')
|
||
context.options |= ssl.OP_NO_SSLv2
|
||
context.options |= ssl.OP_NO_SSLv3
|
||
|
||
return context
|
||
|
||
@staticmethod
|
||
def create_client_context(
|
||
cert_path: Path = None,
|
||
key_path: Path = None,
|
||
ca_path: Path = None,
|
||
verify: bool = False
|
||
) -> ssl.SSLContext:
|
||
"""
|
||
Создаёт SSL контекст для клиента
|
||
|
||
Args:
|
||
cert_path: Путь к сертификату (для mTLS)
|
||
key_path: Путь к приватному ключу (для mTLS)
|
||
ca_path: Путь к CA сертификату
|
||
verify: Проверять сертификат сервера
|
||
"""
|
||
context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
|
||
context.minimum_version = TLSConfig.TLS_VERSION
|
||
|
||
# Клиентский сертификат (для mTLS)
|
||
if cert_path and key_path:
|
||
context.load_cert_chain(str(cert_path), str(key_path))
|
||
|
||
# Проверка сервера
|
||
if verify and ca_path and ca_path.exists():
|
||
context.load_verify_locations(str(ca_path))
|
||
context.verify_mode = ssl.CERT_REQUIRED
|
||
else:
|
||
context.check_hostname = False
|
||
context.verify_mode = ssl.CERT_NONE
|
||
|
||
return context
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════════════
|
||
# ЗАЩИЩЁННЫЙ КЛИЕНТ
|
||
# ═══════════════════════════════════════════════════════════════════════════════
|
||
|
||
class SecureNodeClient:
|
||
"""
|
||
Клиент для защищённой связи с другими узлами
|
||
"""
|
||
|
||
def __init__(self, cert_manager: CertificateManager = None):
|
||
self.cert_manager = cert_manager or CertificateManager()
|
||
|
||
async def connect(
|
||
self,
|
||
host: str,
|
||
port: int = TLSConfig.SECURE_PORT,
|
||
timeout: float = TLSConfig.CONNECT_TIMEOUT
|
||
) -> Tuple[Optional[asyncio.StreamReader], Optional[asyncio.StreamWriter]]:
|
||
"""
|
||
Устанавливает защищённое соединение с узлом
|
||
|
||
Returns:
|
||
(reader, writer) или (None, None) при ошибке
|
||
"""
|
||
try:
|
||
ssl_context = SecureContext.create_client_context(verify=False)
|
||
|
||
reader, writer = await asyncio.wait_for(
|
||
asyncio.open_connection(
|
||
host, port,
|
||
ssl=ssl_context
|
||
),
|
||
timeout=timeout
|
||
)
|
||
|
||
logger.debug(f"🔐 TLS соединение с {host}:{port}")
|
||
return reader, writer
|
||
|
||
except asyncio.TimeoutError:
|
||
logger.warning(f"⏰ Таймаут подключения к {host}:{port}")
|
||
return None, None
|
||
except Exception as e:
|
||
logger.warning(f"❌ Ошибка TLS подключения к {host}:{port}: {e}")
|
||
return None, None
|
||
|
||
async def send_message(
|
||
self,
|
||
host: str,
|
||
message: Dict[str, Any],
|
||
port: int = TLSConfig.SECURE_PORT
|
||
) -> Optional[Dict[str, Any]]:
|
||
"""
|
||
Отправляет сообщение и получает ответ
|
||
|
||
Args:
|
||
host: IP адрес узла
|
||
message: Сообщение (dict)
|
||
port: Порт
|
||
|
||
Returns:
|
||
Ответ (dict) или None
|
||
"""
|
||
reader, writer = await self.connect(host, port)
|
||
if not reader or not writer:
|
||
return None
|
||
|
||
try:
|
||
# Отправляем
|
||
data = json.dumps(message).encode('utf-8')
|
||
writer.write(len(data).to_bytes(4, 'big'))
|
||
writer.write(data)
|
||
await writer.drain()
|
||
|
||
# Получаем ответ
|
||
length_bytes = await asyncio.wait_for(
|
||
reader.read(4),
|
||
timeout=TLSConfig.READ_TIMEOUT
|
||
)
|
||
if not length_bytes:
|
||
return None
|
||
|
||
length = int.from_bytes(length_bytes, 'big')
|
||
response_data = await asyncio.wait_for(
|
||
reader.read(length),
|
||
timeout=TLSConfig.READ_TIMEOUT
|
||
)
|
||
|
||
return json.loads(response_data.decode('utf-8'))
|
||
|
||
except Exception as e:
|
||
logger.warning(f"❌ Ошибка отправки: {e}")
|
||
return None
|
||
finally:
|
||
writer.close()
|
||
await writer.wait_closed()
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════════════
|
||
# ЗАЩИЩЁННЫЙ СЕРВЕР
|
||
# ═══════════════════════════════════════════════════════════════════════════════
|
||
|
||
class SecureNodeServer:
|
||
"""
|
||
Сервер для приёма защищённых соединений
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
node_name: str,
|
||
node_address: str,
|
||
cert_manager: CertificateManager = None
|
||
):
|
||
self.node_name = node_name
|
||
self.node_address = node_address
|
||
self.cert_manager = cert_manager or CertificateManager()
|
||
|
||
self._server = None
|
||
self._handlers = {}
|
||
|
||
# Проверяем/создаём сертификаты
|
||
self._ensure_certificates()
|
||
|
||
def _ensure_certificates(self):
|
||
"""Убеждаемся что сертификаты существуют"""
|
||
cert_path, key_path = self.cert_manager.get_cert_paths(self.node_name)
|
||
|
||
if not cert_path or not key_path:
|
||
logger.info(f"🔐 Генерация сертификата для {self.node_name}...")
|
||
self.cert_path, self.key_path = self.cert_manager.generate_self_signed_cert(
|
||
self.node_address,
|
||
self.node_name
|
||
)
|
||
else:
|
||
self.cert_path = cert_path
|
||
self.key_path = key_path
|
||
|
||
def register_handler(self, message_type: str, handler):
|
||
"""Регистрирует обработчик для типа сообщений"""
|
||
self._handlers[message_type] = handler
|
||
|
||
async def _handle_client(
|
||
self,
|
||
reader: asyncio.StreamReader,
|
||
writer: asyncio.StreamWriter
|
||
):
|
||
"""Обрабатывает клиентское соединение"""
|
||
peer = writer.get_extra_info('peername')
|
||
logger.debug(f"🔐 TLS клиент подключился: {peer}")
|
||
|
||
try:
|
||
while True:
|
||
# Читаем длину сообщения
|
||
length_bytes = await asyncio.wait_for(
|
||
reader.read(4),
|
||
timeout=TLSConfig.READ_TIMEOUT
|
||
)
|
||
if not length_bytes:
|
||
break
|
||
|
||
length = int.from_bytes(length_bytes, 'big')
|
||
|
||
# Читаем сообщение
|
||
data = await asyncio.wait_for(
|
||
reader.read(length),
|
||
timeout=TLSConfig.READ_TIMEOUT
|
||
)
|
||
|
||
message = json.loads(data.decode('utf-8'))
|
||
msg_type = message.get("type", "unknown")
|
||
|
||
# Обрабатываем
|
||
handler = self._handlers.get(msg_type, self._default_handler)
|
||
response = await handler(message, peer)
|
||
|
||
# Отправляем ответ
|
||
response_data = json.dumps(response).encode('utf-8')
|
||
writer.write(len(response_data).to_bytes(4, 'big'))
|
||
writer.write(response_data)
|
||
await writer.drain()
|
||
|
||
except asyncio.TimeoutError:
|
||
logger.debug(f"⏰ Клиент {peer} таймаут")
|
||
except Exception as e:
|
||
logger.warning(f"❌ Ошибка обработки клиента {peer}: {e}")
|
||
finally:
|
||
writer.close()
|
||
await writer.wait_closed()
|
||
|
||
async def _default_handler(self, message: Dict, peer) -> Dict:
|
||
"""Обработчик по умолчанию"""
|
||
return {
|
||
"type": "ack",
|
||
"node": self.node_name,
|
||
"address": self.node_address,
|
||
"timestamp": datetime.now(timezone.utc).isoformat()
|
||
}
|
||
|
||
async def start(
|
||
self,
|
||
host: str = "0.0.0.0",
|
||
port: int = TLSConfig.SECURE_PORT
|
||
):
|
||
"""Запускает TLS сервер"""
|
||
ssl_context = SecureContext.create_server_context(
|
||
self.cert_path,
|
||
self.key_path
|
||
)
|
||
|
||
self._server = await asyncio.start_server(
|
||
self._handle_client,
|
||
host, port,
|
||
ssl=ssl_context
|
||
)
|
||
|
||
logger.info(f"🔐 TLS сервер запущен на {host}:{port}")
|
||
logger.info(f" Сертификат: {self.cert_path}")
|
||
|
||
async with self._server:
|
||
await self._server.serve_forever()
|
||
|
||
def stop(self):
|
||
"""Останавливает сервер"""
|
||
if self._server:
|
||
self._server.close()
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════════════
|
||
# HYBRID TLS + ML-KEM-768 (POST-QUANTUM)
|
||
# ═══════════════════════════════════════════════════════════════════════════════
|
||
|
||
class HybridSecureClient:
|
||
"""
|
||
Клиент с гибридной защитой: TLS 1.3 + ML-KEM-768
|
||
|
||
Последовательность:
|
||
1. Устанавливается TLS 1.3 соединение
|
||
2. Внутри TLS выполняется ML-KEM-768 handshake
|
||
3. Сообщения дважды шифруются: TLS + AES-256-GCM(ML-KEM shared secret)
|
||
|
||
Это обеспечивает защиту даже если:
|
||
- TLS взломан квантовым компьютером
|
||
- ML-KEM-768 имеет уязвимость
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
kem_manager: 'NodeKEMManager' = None,
|
||
cert_manager: CertificateManager = None
|
||
):
|
||
self.cert_manager = cert_manager or CertificateManager()
|
||
self.kem = kem_manager
|
||
self._tls_client = SecureNodeClient(self.cert_manager)
|
||
|
||
async def establish_hybrid_channel(
|
||
self,
|
||
host: str,
|
||
port: int = TLSConfig.SECURE_PORT,
|
||
timeout: float = TLSConfig.CONNECT_TIMEOUT
|
||
) -> Optional[Tuple[asyncio.StreamReader, asyncio.StreamWriter, 'PQSecureChannel']]:
|
||
"""
|
||
Устанавливает гибридный защищённый канал.
|
||
|
||
Returns:
|
||
(reader, writer, pq_channel) или None
|
||
"""
|
||
if not HAS_PQ or not self.kem:
|
||
logger.warning("Post-quantum not available, using TLS only")
|
||
reader, writer = await self._tls_client.connect(host, port, timeout)
|
||
return (reader, writer, None) if reader else None
|
||
|
||
try:
|
||
# 1. Устанавливаем TLS соединение
|
||
ssl_context = SecureContext.create_client_context(verify=False)
|
||
reader, writer = await asyncio.wait_for(
|
||
asyncio.open_connection(host, port, ssl=ssl_context),
|
||
timeout=timeout
|
||
)
|
||
|
||
# 2. Отправляем ML-KEM-768 handshake внутри TLS
|
||
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()
|
||
|
||
# 3. Получаем ответ
|
||
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())
|
||
|
||
# 4. Обрабатываем ML-KEM-768 ответ
|
||
if response.get("type") == "kem_response":
|
||
pq_channel = self.kem.process_handshake_response(response)
|
||
if pq_channel:
|
||
logger.info(f"🔐 Hybrid channel established with {host}")
|
||
return (reader, writer, pq_channel)
|
||
|
||
logger.warning(f"PQ handshake failed with {host}")
|
||
writer.close()
|
||
await writer.wait_closed()
|
||
return None
|
||
|
||
except Exception as e:
|
||
logger.error(f"Hybrid channel error: {e}")
|
||
return None
|
||
|
||
async def send_pq_message(
|
||
self,
|
||
reader: asyncio.StreamReader,
|
||
writer: asyncio.StreamWriter,
|
||
pq_channel: 'PQSecureChannel',
|
||
message: Dict[str, Any]
|
||
) -> Optional[Dict[str, Any]]:
|
||
"""
|
||
Отправляет сообщение с двойным шифрованием.
|
||
|
||
Шифрование: TLS(ML-KEM-768/AES-256-GCM(message))
|
||
"""
|
||
try:
|
||
# 1. Сериализуем и шифруем ML-KEM-768
|
||
plain_data = json.dumps(message).encode()
|
||
encrypted = pq_channel.encrypt(plain_data)
|
||
if not encrypted:
|
||
return None
|
||
|
||
# 2. Отправляем через TLS
|
||
writer.write(len(encrypted).to_bytes(4, 'big'))
|
||
writer.write(encrypted)
|
||
await writer.drain()
|
||
|
||
# 3. Получаем зашифрованный ответ
|
||
length_bytes = await asyncio.wait_for(
|
||
reader.read(4),
|
||
timeout=TLSConfig.READ_TIMEOUT
|
||
)
|
||
length = int.from_bytes(length_bytes, 'big')
|
||
encrypted_response = await asyncio.wait_for(
|
||
reader.read(length),
|
||
timeout=TLSConfig.READ_TIMEOUT
|
||
)
|
||
|
||
# 4. Расшифровываем ML-KEM-768
|
||
decrypted = pq_channel.decrypt(encrypted_response)
|
||
if not decrypted:
|
||
return None
|
||
|
||
return json.loads(decrypted.decode())
|
||
|
||
except Exception as e:
|
||
logger.error(f"PQ message error: {e}")
|
||
return None
|
||
|
||
|
||
class HybridSecureServer:
|
||
"""
|
||
Сервер с гибридной защитой: TLS 1.3 + ML-KEM-768
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
node_name: str,
|
||
node_address: str,
|
||
kem_manager: 'NodeKEMManager' = None,
|
||
cert_manager: CertificateManager = None
|
||
):
|
||
self.node_name = node_name
|
||
self.node_address = node_address
|
||
self.kem = kem_manager
|
||
self.cert_manager = cert_manager or CertificateManager()
|
||
|
||
self._server = None
|
||
self._handlers = {}
|
||
self._pq_channels: Dict[str, 'PQSecureChannel'] = {}
|
||
|
||
# Убеждаемся что сертификаты существуют
|
||
self._ensure_certificates()
|
||
|
||
def _ensure_certificates(self):
|
||
"""Генерируем сертификаты если нужно"""
|
||
cert_path, key_path = self.cert_manager.get_cert_paths(self.node_name)
|
||
if not cert_path:
|
||
self.cert_path, self.key_path = self.cert_manager.generate_self_signed_cert(
|
||
self.node_address, self.node_name
|
||
)
|
||
else:
|
||
self.cert_path = cert_path
|
||
self.key_path = key_path
|
||
|
||
def register_handler(self, message_type: str, handler):
|
||
"""Регистрирует обработчик для типа сообщений"""
|
||
self._handlers[message_type] = handler
|
||
|
||
async def _handle_client(
|
||
self,
|
||
reader: asyncio.StreamReader,
|
||
writer: asyncio.StreamWriter
|
||
):
|
||
"""Обрабатывает клиентское соединение с hybrid security"""
|
||
peer = writer.get_extra_info('peername')
|
||
pq_channel = None
|
||
|
||
try:
|
||
while True:
|
||
# Читаем сообщение
|
||
length_bytes = await asyncio.wait_for(
|
||
reader.read(4),
|
||
timeout=TLSConfig.READ_TIMEOUT
|
||
)
|
||
if not length_bytes:
|
||
break
|
||
|
||
length = int.from_bytes(length_bytes, 'big')
|
||
data = await asyncio.wait_for(
|
||
reader.read(length),
|
||
timeout=TLSConfig.READ_TIMEOUT
|
||
)
|
||
|
||
# Проверяем: это ML-KEM handshake или зашифрованное сообщение?
|
||
if pq_channel is None:
|
||
# Первое сообщение — пробуем как JSON handshake
|
||
try:
|
||
message = json.loads(data.decode())
|
||
|
||
if message.get("type") == "kem_handshake" and self.kem:
|
||
# ML-KEM-768 handshake
|
||
result = self.kem.process_handshake_request(message)
|
||
if result:
|
||
response, pq_channel = result
|
||
self._pq_channels[message["node_address"]] = pq_channel
|
||
|
||
response_data = json.dumps(response).encode()
|
||
writer.write(len(response_data).to_bytes(4, 'big'))
|
||
writer.write(response_data)
|
||
await writer.drain()
|
||
logger.info(f"🔐 PQ handshake complete with {peer}")
|
||
continue
|
||
else:
|
||
# Обычное TLS сообщение
|
||
response = await self._process_message(message, peer)
|
||
response_data = json.dumps(response).encode()
|
||
writer.write(len(response_data).to_bytes(4, 'big'))
|
||
writer.write(response_data)
|
||
await writer.drain()
|
||
continue
|
||
|
||
except json.JSONDecodeError:
|
||
# Не JSON — пропускаем
|
||
pass
|
||
else:
|
||
# PQ канал установлен — расшифровываем
|
||
decrypted = pq_channel.decrypt(data)
|
||
if decrypted:
|
||
message = json.loads(decrypted.decode())
|
||
response = await self._process_message(message, peer)
|
||
|
||
# Шифруем ответ
|
||
response_data = json.dumps(response).encode()
|
||
encrypted = pq_channel.encrypt(response_data)
|
||
if encrypted:
|
||
writer.write(len(encrypted).to_bytes(4, 'big'))
|
||
writer.write(encrypted)
|
||
await writer.drain()
|
||
|
||
except asyncio.TimeoutError:
|
||
pass
|
||
except Exception as e:
|
||
logger.warning(f"Hybrid client error {peer}: {e}")
|
||
finally:
|
||
writer.close()
|
||
await writer.wait_closed()
|
||
|
||
async def _process_message(self, message: Dict, peer) -> Dict:
|
||
"""Обрабатывает сообщение"""
|
||
msg_type = message.get("type", "unknown")
|
||
handler = self._handlers.get(msg_type, self._default_handler)
|
||
return await handler(message, peer)
|
||
|
||
async def _default_handler(self, message: Dict, peer) -> Dict:
|
||
"""Обработчик по умолчанию"""
|
||
return {
|
||
"type": "ack",
|
||
"node": self.node_name,
|
||
"address": self.node_address,
|
||
"pq_enabled": self.kem is not None,
|
||
"timestamp": datetime.now(timezone.utc).isoformat()
|
||
}
|
||
|
||
async def start(self, host: str = "0.0.0.0", port: int = TLSConfig.SECURE_PORT):
|
||
"""Запускает hybrid TLS сервер"""
|
||
ssl_context = SecureContext.create_server_context(self.cert_path, self.key_path)
|
||
|
||
self._server = await asyncio.start_server(
|
||
self._handle_client,
|
||
host, port,
|
||
ssl=ssl_context
|
||
)
|
||
|
||
pq_status = "ML-KEM-768 enabled" if self.kem else "TLS only"
|
||
logger.info(f"🔐 Hybrid server started on {host}:{port} ({pq_status})")
|
||
|
||
async with self._server:
|
||
await self._server.serve_forever()
|
||
|
||
def stop(self):
|
||
"""Останавливает сервер"""
|
||
if self._server:
|
||
self._server.close()
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════════════
|
||
# CLI
|
||
# ═══════════════════════════════════════════════════════════════════════════════
|
||
|
||
if __name__ == "__main__":
|
||
import sys
|
||
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||
)
|
||
|
||
print("🔐 Montana Node TLS")
|
||
print("=" * 50)
|
||
|
||
if len(sys.argv) > 1:
|
||
cmd = sys.argv[1]
|
||
|
||
if cmd == "gen-cert":
|
||
node_name = sys.argv[2] if len(sys.argv) > 2 else "test"
|
||
node_address = sys.argv[3] if len(sys.argv) > 3 else "mt0000000000000000000000000000000000000000"
|
||
|
||
cm = CertificateManager()
|
||
cert, key = cm.generate_self_signed_cert(node_address, node_name)
|
||
print(f"✅ Сертификат: {cert}")
|
||
print(f"✅ Ключ: {key}")
|
||
|
||
elif cmd == "server":
|
||
node_name = sys.argv[2] if len(sys.argv) > 2 else "test"
|
||
node_address = "mt0000000000000000000000000000000000000000"
|
||
|
||
server = SecureNodeServer(node_name, node_address)
|
||
|
||
print(f"🚀 Запуск TLS сервера {node_name}...")
|
||
try:
|
||
asyncio.run(server.start())
|
||
except KeyboardInterrupt:
|
||
print("\n⏹ Остановлено")
|
||
|
||
elif cmd == "test":
|
||
host = sys.argv[2] if len(sys.argv) > 2 else "127.0.0.1"
|
||
|
||
async def test_connection():
|
||
client = SecureNodeClient()
|
||
response = await client.send_message(host, {"type": "ping"})
|
||
if response:
|
||
print(f"✅ Ответ: {response}")
|
||
else:
|
||
print("❌ Нет ответа")
|
||
|
||
asyncio.run(test_connection())
|
||
|
||
else:
|
||
print(f"Неизвестная команда: {cmd}")
|
||
else:
|
||
print("""
|
||
Использование:
|
||
python node_tls.py gen-cert <node_name> [address] — генерация сертификата
|
||
python node_tls.py server <node_name> — запуск TLS сервера
|
||
python node_tls.py test <host> — тест подключения
|
||
""")
|