memory/encryption.py

from __future__ import annotations

import base64
import json
import os
from typing import Any, Dict, Optional

from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF


_INFO = b"inferenceport-ai-shield-storage-v1"


def _secret() -> Optional[bytes]:
    raw = os.getenv("SHIELD_ENCRYPTION_SECRET") or os.getenv("LIGHTNING_TRANSFER_TOKEN")
    if not raw:
        return None
    return raw.encode("utf-8")


def derive_storage_key() -> Optional[bytes]:
    secret = _secret()
    if not secret:
        return None
    hkdf = HKDF(
        algorithm=hashes.SHA256(),
        length=32,
        salt=None,
        info=_INFO,
    )
    return hkdf.derive(secret)


def encrypt_json(data: Dict[str, Any]) -> Optional[Dict[str, str]]:
    key = derive_storage_key()
    if not key:
        return None
    nonce = os.urandom(12)
    plaintext = json.dumps(data, separators=(",", ":"), sort_keys=True).encode("utf-8")
    ciphertext = AESGCM(key).encrypt(nonce, plaintext, None)
    return {
        "nonce": base64.urlsafe_b64encode(nonce).decode("ascii"),
        "ciphertext": base64.urlsafe_b64encode(ciphertext).decode("ascii"),
    }


def decrypt_json(envelope: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    key = derive_storage_key()
    if not key:
        return None
    try:
        nonce = base64.urlsafe_b64decode(str(envelope["nonce"]).encode("ascii"))
        ciphertext = base64.urlsafe_b64decode(str(envelope["ciphertext"]).encode("ascii"))
        plaintext = AESGCM(key).decrypt(nonce, ciphertext, None)
        data = json.loads(plaintext.decode("utf-8"))
        return data if isinstance(data, dict) else None
    except Exception:
        return None