Trixy3/trixy_core/network/protocol.py

"""
Trixy Protocol Implementation

This module implements the custom Trixy binary streaming protocol as specified
in CLAUDE.md. It provides:

- Binary message structure with magic number "TRXI"
- Version fields, datetime, 32-bit option flags, MD5 checksum
- Support for serialized Python classes and dictionaries
- Compression, encryption, and encoding options
- Hard-coded performance commands
- Thread-safe serialization/deserialization

Message Structure:
    [4 bytes] Magic Number: "TRXI"
    [4 bytes] Version Major (int)
    [4 bytes] Version Minor (int) 
    [4 bytes] Version Revision (int)
    [8 bytes] DateTime (binary)
    [4 bytes] Options (32-bit flags)
    [16 bytes] MD5 Checksum
    [4 bytes] Class Name Length
    [variable] Class Name
    [4 bytes] Serialized Data Length
    [variable] Serialized Class Instance

Option Flags (32-bit):
    Bit 1: GZ compressed
    Bit 2: Encrypted
    Bit 3: JSON format (else binary)
    Bit 4: Requires "received" response
    Bit 5: Base64 encoded
    Bit 6: Multi-part message
    Bit 7: Dictionary instead of class
    Bit 8: Silent (no debug logging)
"""

import struct
import time
import hashlib
import gzip
import json
import pickle
import base64
import threading
from typing import Union, Dict, Any, Optional, Tuple, Type
from dataclasses import dataclass
from enum import IntFlag, Enum
from datetime import datetime
import io

# Import encryption if available
try:
    from cryptography.fernet import Fernet
    ENCRYPTION_AVAILABLE = True
except ImportError:
    ENCRYPTION_AVAILABLE = False


def pprint(message: str) -> None:
    """Protocol logging function."""
    print(f"[PROTOCOL] {message}")


# Protocol constants
MAGIC_NUMBER = b"TRXI"
HEADER_SIZE = 48  # Fixed header size in bytes


class MessageFlags(IntFlag):
    """32-bit option flags for Trixy protocol messages."""
    NONE = 0
    GZ_COMPRESSED = 1 << 0       # Bit 1: GZ compressed
    ENCRYPTED = 1 << 1           # Bit 2: Encrypted
    JSON_FORMAT = 1 << 2         # Bit 3: JSON format (else binary)
    REQUIRES_ACK = 1 << 3        # Bit 4: Requires "received" response
    BASE64_ENCODED = 1 << 4      # Bit 5: Base64 encoded
    MULTI_PART = 1 << 5          # Bit 6: Multi-part message
    DICTIONARY = 1 << 6          # Bit 7: Dictionary instead of class
    SILENT = 1 << 7              # Bit 8: Silent (no debug logging)


@dataclass
class ProtocolVersion:
    """Protocol version information."""
    major: int
    minor: int
    revision: int
    
    def __str__(self) -> str:
        return f"{self.major}.{self.minor}.{self.revision}"
    
    def to_bytes(self) -> bytes:
        """Convert version to bytes for protocol header."""
        return struct.pack(">III", self.major, self.minor, self.revision)
    
    @classmethod
    def from_bytes(cls, data: bytes) -> 'ProtocolVersion':
        """Create version from bytes."""
        major, minor, revision = struct.unpack(">III", data)
        return cls(major, minor, revision)


@dataclass
class MessageOptions:
    """Message options and metadata."""
    flags: MessageFlags = MessageFlags.NONE
    timestamp: Optional[datetime] = None
    encryption_key: Optional[bytes] = None
    part_number: int = 0
    total_parts: int = 1
    
    def __post_init__(self):
        if self.timestamp is None:
            self.timestamp = datetime.now()


class ProtocolError(Exception):
    """Base exception for protocol errors."""
    pass


class InvalidMessageError(ProtocolError):
    """Raised when message format is invalid."""
    pass


class ChecksumMismatchError(ProtocolError):
    """Raised when MD5 checksum doesn't match."""
    pass


class UnsupportedVersionError(ProtocolError):
    """Raised when protocol version is unsupported."""
    pass


class CompressionError(ProtocolError):
    """Raised when compression/decompression fails."""
    pass


class EncryptionError(ProtocolError):
    """Raised when encryption/decryption fails."""
    pass


class TrixyMessage:
    """
    Represents a complete Trixy protocol message.
    
    This class handles the binary message format including all headers,
    options, and payload data according to the protocol specification.
    """
    
    def __init__(
        self,
        class_name: str,
        data: Any,
        version: Optional[ProtocolVersion] = None,
        options: Optional[MessageOptions] = None
    ):
        """
        Initialize a Trixy message.
        
        Args:
            class_name: Name of the command class or "HARDCODED" for hard-coded commands
            data: The data to serialize (class instance, dict, or string)
            version: Protocol version (uses current if None)
            options: Message options (uses defaults if None)
        """
        self.class_name = class_name
        self.data = data
        self.version = version or ProtocolVersion(1, 0, 0)
        self.options = options or MessageOptions()
        
        # Calculated fields
        self.checksum: Optional[bytes] = None
        self.serialized_data: Optional[bytes] = None
        self._lock = threading.RLock()
    
    def serialize(self) -> bytes:
        """
        Serialize the message to binary format according to protocol specification.
        
        Returns:
            bytes: Complete binary message
            
        Raises:
            ProtocolError: If serialization fails
        """
        with self._lock:
            try:
                # Serialize the payload data
                if self.class_name == "HARDCODED":
                    # Hard-coded commands are sent as UTF-8 strings
                    payload_data = str(self.data).encode('utf-8')
                elif self.options.flags & MessageFlags.DICTIONARY:
                    # Serialize as dictionary
                    if self.options.flags & MessageFlags.JSON_FORMAT:
                        payload_data = json.dumps(self.data).encode('utf-8')
                    else:
                        payload_data = pickle.dumps(self.data)
                else:
                    # Serialize as class instance
                    if self.options.flags & MessageFlags.JSON_FORMAT:
                        # Convert class to dict then JSON
                        if hasattr(self.data, '__dict__'):
                            class_dict = self.data.__dict__.copy()
                            class_dict['__class__'] = self.data.__class__.__name__
                            payload_data = json.dumps(class_dict).encode('utf-8')
                        else:
                            payload_data = json.dumps(str(self.data)).encode('utf-8')
                    else:
                        payload_data = pickle.dumps(self.data)
                
                # Apply compression if requested
                if self.options.flags & MessageFlags.GZ_COMPRESSED:
                    payload_data = compress_data(payload_data)
                
                # Apply encryption if requested
                if self.options.flags & MessageFlags.ENCRYPTED:
                    if not self.options.encryption_key:
                        raise EncryptionError("Encryption requested but no key provided")
                    payload_data = encrypt_data(payload_data, self.options.encryption_key)
                
                # Apply base64 encoding if requested
                if self.options.flags & MessageFlags.BASE64_ENCODED:
                    payload_data = base64.b64encode(payload_data)
                
                # Store serialized data for checksum calculation
                self.serialized_data = payload_data
                
                # Build the complete message
                message = self._build_message(payload_data)
                
                if not (self.options.flags & MessageFlags.SILENT):
                    pprint(f"Serialized message: {self.class_name} ({len(message)} bytes)")
                
                return message
                
            except Exception as e:
                raise ProtocolError(f"Failed to serialize message: {e}") from e
    
    def _build_message(self, payload_data: bytes) -> bytes:
        """Build the complete binary message with headers."""
        # Encode class name
        class_name_bytes = self.class_name.encode('utf-8')
        
        # Build header components
        magic = MAGIC_NUMBER
        version_bytes = self.version.to_bytes()
        timestamp_bytes = struct.pack(">d", self.options.timestamp.timestamp())
        flags_bytes = struct.pack(">I", int(self.options.flags))
        
        # Calculate MD5 checksum of class name + payload
        checksum_data = class_name_bytes + payload_data
        checksum = calculate_checksum(checksum_data)
        
        # Length fields
        class_name_length = struct.pack(">I", len(class_name_bytes))
        payload_length = struct.pack(">I", len(payload_data))
        
        # Assemble complete message
        message = (
            magic +                    # 4 bytes
            version_bytes +            # 12 bytes (3 * 4)
            timestamp_bytes +          # 8 bytes
            flags_bytes +              # 4 bytes
            checksum +                 # 16 bytes
            class_name_length +        # 4 bytes
            class_name_bytes +         # variable
            payload_length +           # 4 bytes
            payload_data               # variable
        )
        
        return message
    
    @classmethod
    def deserialize(cls, data: bytes) -> 'TrixyMessage':
        """
        Deserialize binary data into a TrixyMessage.
        
        Args:
            data: Binary message data
            
        Returns:
            TrixyMessage: Deserialized message
            
        Raises:
            ProtocolError: If deserialization fails
        """
        try:
            if len(data) < HEADER_SIZE:
                raise InvalidMessageError(f"Message too short: {len(data)} bytes")
            
            # Parse header
            offset = 0
            
            # Magic number
            magic = data[offset:offset + 4]
            if magic != MAGIC_NUMBER:
                raise InvalidMessageError(f"Invalid magic number: {magic}")
            offset += 4
            
            # Version
            version = ProtocolVersion.from_bytes(data[offset:offset + 12])
            offset += 12
            
            # Timestamp
            timestamp_float = struct.unpack(">d", data[offset:offset + 8])[0]
            timestamp = datetime.fromtimestamp(timestamp_float)
            offset += 8
            
            # Flags
            flags = MessageFlags(struct.unpack(">I", data[offset:offset + 4])[0])
            offset += 4
            
            # Checksum
            received_checksum = data[offset:offset + 16]
            offset += 16
            
            # Class name length and data
            class_name_length = struct.unpack(">I", data[offset:offset + 4])[0]
            offset += 4
            
            class_name = data[offset:offset + class_name_length].decode('utf-8')
            offset += class_name_length
            
            # Payload length and data
            payload_length = struct.unpack(">I", data[offset:offset + 4])[0]
            offset += 4
            
            payload_data = data[offset:offset + payload_length]
            
            # Verify checksum
            checksum_data = class_name.encode('utf-8') + payload_data
            calculated_checksum = calculate_checksum(checksum_data)
            
            if received_checksum != calculated_checksum:
                raise ChecksumMismatchError("MD5 checksum mismatch")
            
            # Process payload data
            processed_data = cls._process_payload(payload_data, flags, class_name)
            
            # Create message options
            options = MessageOptions(flags=flags, timestamp=timestamp)
            
            # Create message instance
            message = cls(class_name, processed_data, version, options)
            message.checksum = received_checksum
            message.serialized_data = payload_data
            
            if not (flags & MessageFlags.SILENT):
                pprint(f"Deserialized message: {class_name} ({len(data)} bytes)")
            
            return message
            
        except Exception as e:
            if isinstance(e, ProtocolError):
                raise
            raise ProtocolError(f"Failed to deserialize message: {e}") from e
    
    @classmethod
    def _process_payload(cls, payload_data: bytes, flags: MessageFlags, class_name: str) -> Any:
        """Process payload data according to flags."""
        data = payload_data
        
        # Reverse base64 encoding
        if flags & MessageFlags.BASE64_ENCODED:
            data = base64.b64decode(data)
        
        # Reverse encryption (would need key - placeholder for now)
        if flags & MessageFlags.ENCRYPTED:
            raise EncryptionError("Cannot decrypt without key")
        
        # Reverse compression
        if flags & MessageFlags.GZ_COMPRESSED:
            data = decompress_data(data)
        
        # Special handling for hard-coded commands
        if class_name == "HARDCODED":
            return data.decode('utf-8')
        
        # Deserialize based on format
        if flags & MessageFlags.DICTIONARY:
            if flags & MessageFlags.JSON_FORMAT:
                return json.loads(data.decode('utf-8'))
            else:
                return pickle.loads(data)
        else:
            if flags & MessageFlags.JSON_FORMAT:
                json_data = json.loads(data.decode('utf-8'))
                # If it's a class dict, we'd need to reconstruct the class
                # For now just return the dict
                return json_data
            else:
                return pickle.loads(data)


class TrixyProtocol:
    """
    Main protocol handler for Trixy network communication.
    
    This class provides high-level methods for creating, sending, and receiving
    Trixy protocol messages. It handles all serialization, compression, and
    other protocol features.
    """
    
    def __init__(self, version: Optional[ProtocolVersion] = None):
        """
        Initialize the protocol handler.
        
        Args:
            version: Protocol version to use (defaults to 1.0.0)
        """
        self.version = version or ProtocolVersion(1, 0, 0)
        self._lock = threading.RLock()
        self._encryption_keys: Dict[str, bytes] = {}
        
        pprint(f"Trixy protocol initialized (version {self.version})")
    
    def create_message(
        self,
        class_name: str,
        data: Any,
        options: Optional[MessageOptions] = None
    ) -> TrixyMessage:
        """
        Create a new protocol message.
        
        Args:
            class_name: Name of the command class
            data: Data to serialize
            options: Message options
            
        Returns:
            TrixyMessage: Created message
        """
        with self._lock:
            return TrixyMessage(class_name, data, self.version, options)
    
    def create_hardcoded_message(self, command: str) -> TrixyMessage:
        """
        Create a hard-coded command message for performance.
        
        Args:
            command: Hard-coded command string (TRXINOOP, TRXIPING, etc.)
            
        Returns:
            TrixyMessage: Created message
        """
        options = MessageOptions(flags=MessageFlags.SILENT)  # Hard-coded commands are silent
        return self.create_message("HARDCODED", command, options)
    
    def serialize_message(self, message: TrixyMessage) -> bytes:
        """
        Serialize a message to binary format.
        
        Args:
            message: Message to serialize
            
        Returns:
            bytes: Serialized message
        """
        return message.serialize()
    
    def deserialize_message(self, data: bytes) -> TrixyMessage:
        """
        Deserialize binary data to a message.
        
        Args:
            data: Binary message data
            
        Returns:
            TrixyMessage: Deserialized message
        """
        return TrixyMessage.deserialize(data)
    
    def is_hardcoded_command(self, data: bytes) -> bool:
        """
        Check if data starts with a hard-coded command.
        
        Args:
            data: Raw data from socket
            
        Returns:
            bool: True if it's a hard-coded command
        """
        try:
            # Hard-coded commands are simple text strings
            text = data.decode('utf-8', errors='ignore')
            return any(text.startswith(cmd) for cmd in ["TRXI"])
        except:
            return False
    
    def parse_hardcoded_command(self, data: bytes) -> Tuple[str, str]:
        """
        Parse a hard-coded command from raw data.
        
        Args:
            data: Raw command data
            
        Returns:
            Tuple[str, str]: (command, arguments)
        """
        try:
            text = data.decode('utf-8').strip()
            parts = text.split(' ', 1)
            command = parts[0]
            args = parts[1] if len(parts) > 1 else ""
            return command, args
        except Exception as e:
            raise ProtocolError(f"Failed to parse hard-coded command: {e}") from e
    
    def set_encryption_key(self, key_id: str, key: bytes) -> None:
        """
        Set an encryption key for encrypted messages.
        
        Args:
            key_id: Identifier for the key
            key: Encryption key bytes
        """
        with self._lock:
            self._encryption_keys[key_id] = key
            pprint(f"Encryption key set for: {key_id}")
    
    def get_encryption_key(self, key_id: str) -> Optional[bytes]:
        """
        Get an encryption key by ID.
        
        Args:
            key_id: Key identifier
            
        Returns:
            Optional[bytes]: Encryption key or None if not found
        """
        with self._lock:
            return self._encryption_keys.get(key_id)


# Utility functions

def calculate_checksum(data: bytes) -> bytes:
    """
    Calculate MD5 checksum for data.
    
    Args:
        data: Data to checksum
        
    Returns:
        bytes: 16-byte MD5 hash
    """
    return hashlib.md5(data).digest()


def compress_data(data: bytes) -> bytes:
    """
    Compress data using gzip.
    
    Args:
        data: Data to compress
        
    Returns:
        bytes: Compressed data
        
    Raises:
        CompressionError: If compression fails
    """
    try:
        return gzip.compress(data)
    except Exception as e:
        raise CompressionError(f"Failed to compress data: {e}") from e


def decompress_data(data: bytes) -> bytes:
    """
    Decompress gzip data.
    
    Args:
        data: Compressed data
        
    Returns:
        bytes: Decompressed data
        
    Raises:
        CompressionError: If decompression fails
    """
    try:
        return gzip.decompress(data)
    except Exception as e:
        raise CompressionError(f"Failed to decompress data: {e}") from e


def encrypt_data(data: bytes, key: bytes) -> bytes:
    """
    Encrypt data using Fernet symmetric encryption.
    
    Args:
        data: Data to encrypt
        key: Encryption key
        
    Returns:
        bytes: Encrypted data
        
    Raises:
        EncryptionError: If encryption fails or not available
    """
    if not ENCRYPTION_AVAILABLE:
        raise EncryptionError("Cryptography library not available")
    
    try:
        fernet = Fernet(key)
        return fernet.encrypt(data)
    except Exception as e:
        raise EncryptionError(f"Failed to encrypt data: {e}") from e


def decrypt_data(data: bytes, key: bytes) -> bytes:
    """
    Decrypt data using Fernet symmetric encryption.
    
    Args:
        data: Encrypted data
        key: Decryption key
        
    Returns:
        bytes: Decrypted data
        
    Raises:
        EncryptionError: If decryption fails or not available
    """
    if not ENCRYPTION_AVAILABLE:
        raise EncryptionError("Cryptography library not available")
    
    try:
        fernet = Fernet(key)
        return fernet.decrypt(data)
    except Exception as e:
        raise EncryptionError(f"Failed to decrypt data: {e}") from e


def serialize_object(obj: Any, use_json: bool = False) -> bytes:
    """
    Serialize an object to bytes.
    
    Args:
        obj: Object to serialize
        use_json: Use JSON instead of pickle
        
    Returns:
        bytes: Serialized object
    """
    try:
        if use_json:
            if hasattr(obj, '__dict__'):
                obj_dict = obj.__dict__.copy()
                obj_dict['__class__'] = obj.__class__.__name__
                return json.dumps(obj_dict).encode('utf-8')
            else:
                return json.dumps(obj).encode('utf-8')
        else:
            return pickle.dumps(obj)
    except Exception as e:
        raise ProtocolError(f"Failed to serialize object: {e}") from e


def deserialize_object(data: bytes, use_json: bool = False) -> Any:
    """
    Deserialize bytes to an object.
    
    Args:
        data: Serialized data
        use_json: Expect JSON format
        
    Returns:
        Any: Deserialized object
    """
    try:
        if use_json:
            return json.loads(data.decode('utf-8'))
        else:
            return pickle.loads(data)
    except Exception as e:
        raise ProtocolError(f"Failed to deserialize object: {e}") from e


def generate_encryption_key() -> bytes:
    """
    Generate a new Fernet encryption key.
    
    Returns:
        bytes: Generated key
        
    Raises:
        EncryptionError: If cryptography library not available
    """
    if not ENCRYPTION_AVAILABLE:
        raise EncryptionError("Cryptography library not available")
    
    return Fernet.generate_key()


# Module exports
__all__ = [
    "TrixyProtocol",
    "TrixyMessage",
    "MessageOptions", 
    "MessageFlags",
    "ProtocolVersion",
    "ProtocolError",
    "InvalidMessageError",
    "ChecksumMismatchError",
    "UnsupportedVersionError",
    "CompressionError",
    "EncryptionError",
    "serialize_object",
    "deserialize_object",
    "calculate_checksum",
    "compress_data",
    "decompress_data",
    "encrypt_data",
    "decrypt_data",
    "generate_encryption_key",
    "MAGIC_NUMBER",
    "HEADER_SIZE",
]