Trixy7/trixy_core/network/encryption.py

# -*- coding: utf-8 -*-
"""
Verschlüsselung für das Trixy-Protokoll.

Verwendet AES-256-GCM für sichere Kommunikation.
"""

import os
import secrets
from pathlib import Path
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

from trixy_core.utils.debug import pinfo, pdebug, perror


class TrixyEncryption:
    """
    AES-256-GCM Verschlüsselung für Trixy-Kommunikation.

    Struktur verschlüsselter Daten:
    - Nonce (12 bytes)
    - Ciphertext (variable)
    - Tag (16 bytes, in Ciphertext enthalten)
    """

    NONCE_SIZE = 12
    KEY_SIZE = 32  # 256 bit

    def __init__(self, key: bytes | None = None) -> None:
        """
        Initialisiert die Verschlüsselung.

        Args:
            key: 32-byte AES-Schlüssel, oder None für neuen Schlüssel
        """
        if key is None:
            self._key = self.generate_key()
        else:
            if len(key) != self.KEY_SIZE:
                raise ValueError(f"Schlüssel muss {self.KEY_SIZE} Bytes lang sein")
            self._key = key

        self._aesgcm = AESGCM(self._key)

    @property
    def key(self) -> bytes:
        """Gibt den Schlüssel zurück."""
        return self._key

    @staticmethod
    def generate_key() -> bytes:
        """
        Generiert einen neuen zufälligen Schlüssel.

        Returns:
            32-byte Schlüssel
        """
        return secrets.token_bytes(TrixyEncryption.KEY_SIZE)

    def encrypt(self, plaintext: bytes, associated_data: bytes | None = None) -> bytes:
        """
        Verschlüsselt Daten mit AES-256-GCM.

        Args:
            plaintext: Zu verschlüsselnde Daten
            associated_data: Optionale zusätzliche authentifizierte Daten

        Returns:
            Nonce + Ciphertext (enthält Tag)
        """
        nonce = secrets.token_bytes(self.NONCE_SIZE)
        ciphertext = self._aesgcm.encrypt(nonce, plaintext, associated_data)
        return nonce + ciphertext

    def decrypt(self, ciphertext: bytes, associated_data: bytes | None = None) -> bytes:
        """
        Entschlüsselt Daten mit AES-256-GCM.

        Args:
            ciphertext: Nonce + verschlüsselte Daten
            associated_data: Optionale zusätzliche authentifizierte Daten

        Returns:
            Entschlüsselte Daten

        Raises:
            ValueError: Bei ungültigen Daten oder Authentifizierungsfehler
        """
        if len(ciphertext) < self.NONCE_SIZE + 16:  # 16 = min Tag-Größe
            raise ValueError("Ciphertext zu kurz")

        nonce = ciphertext[:self.NONCE_SIZE]
        encrypted_data = ciphertext[self.NONCE_SIZE:]

        try:
            return self._aesgcm.decrypt(nonce, encrypted_data, associated_data)
        except Exception as e:
            perror(f"Entschlüsselung fehlgeschlagen: {e}")
            raise ValueError("Entschlüsselung fehlgeschlagen") from e

    def save_key(self, path: str | Path) -> None:
        """
        Speichert den Schlüssel in eine Datei.

        Args:
            path: Zielpfad
        """
        file_path = Path(path)
        file_path.parent.mkdir(parents=True, exist_ok=True)

        with open(file_path, "wb") as f:
            f.write(self._key)

        # Nur für Besitzer lesbar
        os.chmod(file_path, 0o600)
        pinfo(f"Verschlüsselungsschlüssel gespeichert: {file_path}")

    @classmethod
    def load_key(cls, path: str | Path) -> "TrixyEncryption":
        """
        Lädt einen Schlüssel aus einer Datei.

        Args:
            path: Dateipfad

        Returns:
            TrixyEncryption-Instanz mit geladenem Schlüssel

        Raises:
            FileNotFoundError: Wenn die Datei nicht existiert
        """
        file_path = Path(path)

        if not file_path.exists():
            raise FileNotFoundError(f"Schlüsseldatei nicht gefunden: {file_path}")

        with open(file_path, "rb") as f:
            key = f.read()

        pinfo(f"Verschlüsselungsschlüssel geladen: {file_path}")
        return cls(key)

    @classmethod
    def load_or_create(cls, path: str | Path) -> "TrixyEncryption":
        """
        Lädt oder erstellt einen Schlüssel.

        Args:
            path: Dateipfad

        Returns:
            TrixyEncryption-Instanz
        """
        file_path = Path(path)

        if file_path.exists():
            return cls.load_key(file_path)
        else:
            instance = cls()
            instance.save_key(file_path)
            return instance