Trixy7/trixy_core/network/protocol.py

# -*- coding: utf-8 -*-
"""
Trixy-Protokoll für die Netzwerkkommunikation.

Binäres Protokoll mit:
- Magic Number 'TRXI'
- Versionierung
- 32-bit Optionsflags
- MD5 Prüfsumme
- Serialisierte Daten
"""

import hashlib
import json
import pickle
import struct
import base64
from dataclasses import dataclass, field
from datetime import datetime
from enum import IntFlag
from typing import Any, Type

from trixy_core.utils.version import PROTOCOL_MAJOR, PROTOCOL_MINOR
from trixy_core.utils.debug import pdebug, perror


# Protokoll-Magic-Number
MAGIC = b"TRXI"
MAGIC_LENGTH = 4

# Hard-coded Befehle für Effizienz
COMMAND_NOOP = b"TRXINOOP"  # Keine Operation
COMMAND_PING = b"TRXIPING"  # Ping-Anfrage
COMMAND_PONG = b"TRXIPONG"  # Ping-Antwort
COMMAND_PRINT = b"TRXIPRNT"  # Debug-Ausgabe
COMMAND_HELLO = b"TRXYHELO"  # Handshake

HARD_COMMANDS = {COMMAND_NOOP, COMMAND_PING, COMMAND_PONG, COMMAND_PRINT, COMMAND_HELLO}


class ProtocolFlags(IntFlag):
    """32-bit Protokoll-Optionen."""
    NONE = 0
    COMPRESSED = 1 << 0       # Daten sind komprimiert (zlib)
    ENCRYPTED = 1 << 1        # Daten sind verschlüsselt (AES-256-GCM)
    JSON = 1 << 2             # Daten sind JSON
    ACK_REQUIRED = 1 << 3     # Bestätigung erforderlich
    BASE64 = 1 << 4           # Daten sind Base64-kodiert
    MULTIPART = 1 << 5        # Teil einer Mehrteiligen Nachricht
    DICT = 1 << 6             # Daten sind ein Dictionary
    SILENT = 1 << 7           # Keine Logging-Ausgabe
    PICKLE = 1 << 8           # Daten sind Pickle-serialisiert
    RESPONSE = 1 << 9         # Ist eine Antwort auf eine Anfrage
    ERROR = 1 << 10           # Enthält Fehlerdaten
    BINARY = 1 << 11          # Rohe Binärdaten


@dataclass
class ProtocolMessage:
    """
    Repräsentiert eine Protokoll-Nachricht.

    Struktur:
    - Magic (4 bytes): 'TRXI'
    - Version (2 bytes): Major.Minor
    - Timestamp (8 bytes): Unix-Timestamp in Millisekunden
    - Flags (4 bytes): 32-bit Optionsflags
    - Checksum (16 bytes): MD5 der Daten
    - Class Name Length (2 bytes): Länge des Klassennamens
    - Class Name (variable): Name der Nachrichtenklasse
    - Data Length (4 bytes): Länge der Daten
    - Data (variable): Serialisierte Daten
    """
    class_name: str = ""
    data: Any = None
    flags: ProtocolFlags = ProtocolFlags.NONE
    timestamp: datetime = field(default_factory=datetime.now)
    checksum: bytes = b""
    raw_data: bytes = field(default=b"", repr=False)

    # Header-Größe ohne variable Teile
    HEADER_SIZE = MAGIC_LENGTH + 2 + 8 + 4 + 16 + 2


class TrixyProtocol:
    """
    Protokoll-Handler für Serialisierung und Deserialisierung.
    """

    def __init__(self, encryption: Any = None) -> None:
        """
        Initialisiert das Protokoll.

        Args:
            encryption: Optionale TrixyEncryption-Instanz
        """
        self._encryption = encryption
        self._registered_classes: dict[str, Type] = {}

    def register_class(self, cls: Type) -> None:
        """
        Registriert eine Klasse für die Deserialisierung.

        Args:
            cls: Zu registrierende Klasse
        """
        self._registered_classes[cls.__name__] = cls

    def set_encryption(self, encryption: Any) -> None:
        """Setzt die Verschlüsselung."""
        self._encryption = encryption

    def serialize(
        self,
        message: Any,
        flags: ProtocolFlags = ProtocolFlags.NONE
    ) -> bytes:
        """
        Serialisiert eine Nachricht ins Protokoll-Format.

        Args:
            message: Zu serialisierende Nachricht
            flags: Protokoll-Flags

        Returns:
            Serialisierte Bytes
        """
        class_name = type(message).__name__
        timestamp = datetime.now()
        timestamp_ms = int(timestamp.timestamp() * 1000)

        # Daten serialisieren
        if flags & ProtocolFlags.JSON:
            if hasattr(message, "__dict__"):
                data = json.dumps(message.__dict__).encode("utf-8")
            else:
                data = json.dumps(message).encode("utf-8")
        elif flags & ProtocolFlags.PICKLE:
            data = pickle.dumps(message)
        elif flags & ProtocolFlags.BINARY:
            data = message if isinstance(message, bytes) else bytes(message)
        else:
            # Standard: Pickle
            flags |= ProtocolFlags.PICKLE
            data = pickle.dumps(message)

        # Base64-Kodierung
        if flags & ProtocolFlags.BASE64:
            data = base64.b64encode(data)

        # Verschlüsselung
        if flags & ProtocolFlags.ENCRYPTED:
            if self._encryption is None:
                perror("Verschlüsselung angefordert, aber nicht konfiguriert")
                flags &= ~ProtocolFlags.ENCRYPTED
            else:
                data = self._encryption.encrypt(data)

        # MD5 Prüfsumme
        checksum = hashlib.md5(data).digest()

        # Klassenname kodieren
        class_name_bytes = class_name.encode("utf-8")

        # Nachricht zusammenbauen
        header = struct.pack(
            ">4sBBQI16sH",
            MAGIC,
            PROTOCOL_MAJOR,
            PROTOCOL_MINOR,
            timestamp_ms,
            int(flags),
            checksum,
            len(class_name_bytes)
        )

        data_length = struct.pack(">I", len(data))

        return header + class_name_bytes + data_length + data

    def deserialize(self, raw_bytes: bytes) -> ProtocolMessage:
        """
        Deserialisiert Bytes in eine Protokoll-Nachricht.

        Args:
            raw_bytes: Rohe Bytes

        Returns:
            Deserialisierte Nachricht

        Raises:
            ValueError: Bei ungültigen Daten
        """
        # Prüfe zuerst auf Hard-coded Befehle (8 Bytes)
        # Diese haben eine andere Struktur als normale Nachrichten
        if len(raw_bytes) >= 8:
            cmd = raw_bytes[:8]
            if cmd in HARD_COMMANDS:
                return ProtocolMessage(
                    class_name=cmd.decode("utf-8"),
                    data=cmd,
                    flags=ProtocolFlags.BINARY
                )

        # Für normale Nachrichten: Header-Größe prüfen
        if len(raw_bytes) < ProtocolMessage.HEADER_SIZE:
            raise ValueError("Nachricht zu kurz für Header")

        # Header parsen
        offset = 0
        magic = raw_bytes[offset:offset + MAGIC_LENGTH]
        offset += MAGIC_LENGTH

        if magic != MAGIC:
            raise ValueError(f"Ungültige Magic Number: {magic}")

        version_major, version_minor = struct.unpack(">BB", raw_bytes[offset:offset + 2])
        offset += 2

        timestamp_ms = struct.unpack(">Q", raw_bytes[offset:offset + 8])[0]
        offset += 8
        timestamp = datetime.fromtimestamp(timestamp_ms / 1000)

        flags = ProtocolFlags(struct.unpack(">I", raw_bytes[offset:offset + 4])[0])
        offset += 4

        checksum = raw_bytes[offset:offset + 16]
        offset += 16

        class_name_length = struct.unpack(">H", raw_bytes[offset:offset + 2])[0]
        offset += 2

        class_name = raw_bytes[offset:offset + class_name_length].decode("utf-8")
        offset += class_name_length

        data_length = struct.unpack(">I", raw_bytes[offset:offset + 4])[0]
        offset += 4

        data = raw_bytes[offset:offset + data_length]

        # Prüfsumme validieren
        calculated_checksum = hashlib.md5(data).digest()
        if checksum != calculated_checksum:
            perror("Prüfsummenfehler bei Protokoll-Nachricht")
            raise ValueError("Prüfsummenfehler")

        # Entschlüsselung
        if flags & ProtocolFlags.ENCRYPTED:
            if self._encryption is None:
                raise ValueError("Verschlüsselte Nachricht, aber keine Verschlüsselung konfiguriert")
            data = self._encryption.decrypt(data)

        # Base64-Dekodierung
        if flags & ProtocolFlags.BASE64:
            data = base64.b64decode(data)

        # Deserialisierung
        deserialized_data: Any = None
        if flags & ProtocolFlags.JSON:
            deserialized_data = json.loads(data.decode("utf-8"))
        elif flags & ProtocolFlags.PICKLE:
            deserialized_data = pickle.loads(data)
        elif flags & ProtocolFlags.BINARY:
            deserialized_data = data
        else:
            deserialized_data = data

        # Klasseninstanz erstellen wenn registriert
        if class_name in self._registered_classes and isinstance(deserialized_data, dict):
            cls = self._registered_classes[class_name]
            try:
                deserialized_data = cls(**deserialized_data)
            except Exception:
                pass  # Behalte Dictionary

        return ProtocolMessage(
            class_name=class_name,
            data=deserialized_data,
            flags=flags,
            timestamp=timestamp,
            checksum=checksum,
            raw_data=raw_bytes
        )

    def create_ping(self) -> bytes:
        """Erstellt eine PING-Nachricht."""
        return COMMAND_PING

    def create_pong(self) -> bytes:
        """Erstellt eine PONG-Nachricht."""
        return COMMAND_PONG

    def create_noop(self) -> bytes:
        """Erstellt eine NOOP-Nachricht."""
        return COMMAND_NOOP

    def create_hello(self) -> bytes:
        """Erstellt eine HELLO-Nachricht."""
        return COMMAND_HELLO

    def is_hard_command(self, data: bytes) -> bool:
        """Prüft, ob die Daten ein Hard-coded Befehl sind."""
        return data[:8] in HARD_COMMANDS if len(data) >= 8 else False

    def get_hard_command(self, data: bytes) -> str | None:
        """Gibt den Hard-coded Befehl zurück oder None."""
        if len(data) >= 8:
            cmd = data[:8]
            if cmd in HARD_COMMANDS:
                return cmd.decode("utf-8")
        return None