cache_layers.py

"""Cache storage layers for TNFR.

This module provides the storage backend implementations for the cache system,
including in-memory, shelve-based, and Redis-based layers. It also handles
secure serialization with HMAC signing to prevent tampering.
"""

from __future__ import annotations

import io
import os
import pickle
import pickletools
import shelve
import threading
import warnings
from abc import ABC, abstractmethod
from collections.abc import Callable, MutableMapping
from typing import Any

from ..errors import TNFRSecurityError, TNFRSecurityWarning, TNFRValueError
from ..security.crypto import create_hmac_signer, create_hmac_validator

__all__ = (
    "CacheLayer",
    "MappingCacheLayer",
    "ShelveCacheLayer",
    "RedisCacheLayer",
    "create_secure_shelve_layer",
    "create_secure_redis_layer",
)

_SIGNATURE_PREFIX = b"TNFRSIG2"
_LEGACY_SIGNATURE_PREFIX = b"TNFRSIG1"
_SIGN_MODE_RAW = 0
_SIGN_MODE_PICKLE = 1
_SIGNATURE_HEADER_SIZE = len(_SIGNATURE_PREFIX) + 1 + 4

# Environment variable to control security warnings for pickle deserialization
_TNFR_ALLOW_UNSIGNED_PICKLE = "TNFR_ALLOW_UNSIGNED_PICKLE"


def create_secure_shelve_layer(
    path: str,
    secret: bytes | str | None = None,
    *,
    flag: str = "c",
    protocol: int | None = None,
    writeback: bool = False,
) -> ShelveCacheLayer:
    """Create a ShelveCacheLayer with HMAC signature validation enabled.

    This is the recommended way to create persistent cache layers that handle
    TNFR structures (EPI, NFR, NetworkX graphs). The outer shelf container is
    decoded without executable globals; version-2 signatures authenticate both
    payload bytes and their interpretation before inner pickle decoding.
    Legacy version-1 signed cache entries must be rebuilt.

    Parameters
    ----------
    path : str
        Path to the shelve database file.
    secret : bytes, str, or None
        Secret key for HMAC signing. If None, reads from TNFR_CACHE_SECRET
        environment variable. In production, **always** set this via environment.
    flag : str, default='c'
        Shelve open flag ('r', 'w', 'c', 'n').
    protocol : int, optional
        Pickle protocol version. Defaults to pickle.HIGHEST_PROTOCOL.
    writeback : bool, default=False
        Enable shelve writeback mode.

    Returns
    -------
    ShelveCacheLayer
        A cache layer with signature validation enabled.

    Raises
    ------
    ValueError
        If no secret is provided and TNFR_CACHE_SECRET is not set.

    Examples
    --------
    >>> # In production, set environment variable:
    >>> # export TNFR_CACHE_SECRET="your-secure-random-key"
    >>>
    >>> layer = create_secure_shelve_layer("coherence.db")
    >>> # Or explicitly provide secret:
    >>> layer = create_secure_shelve_layer("coherence.db", secret=b"my-secret")
    """
    if secret is None:
        secret = os.environ.get("TNFR_CACHE_SECRET")
        if not secret:
            raise TNFRValueError(
                "Secret required for secure cache layer.",
                context={"env_var": "TNFR_CACHE_SECRET"},
                suggestion="set TNFR_CACHE_SECRET environment variable or pass secret parameter.",
            )

    signer = create_hmac_signer(secret)
    validator = create_hmac_validator(secret)

    return ShelveCacheLayer(
        path,
        flag=flag,
        protocol=protocol,
        writeback=writeback,
        signer=signer,
        validator=validator,
        require_signature=True,
    )


def create_secure_redis_layer(
    client: Any | None = None,
    secret: bytes | str | None = None,
    *,
    namespace: str = "tnfr:cache",
    protocol: int | None = None,
) -> RedisCacheLayer:
    """Create a RedisCacheLayer with HMAC signature validation enabled.

    This is the recommended way to create distributed cache layers for TNFR.
    Signature validation protects against arbitrary code execution if Redis
    is compromised or contains tampered data. Version-2 signatures cover the
    raw/pickle mode as well as the payload; version-1 entries must be rebuilt.

    Parameters
    ----------
    client : redis.Redis, optional
        Redis client instance. If None, creates default client.
    secret : bytes, str, or None
        Secret key for HMAC signing. If None, reads from TNFR_CACHE_SECRET
        environment variable.
    namespace : str, default='tnfr:cache'
        Redis key namespace prefix.
    protocol : int, optional
        Pickle protocol version.

    Returns
    -------
    RedisCacheLayer
        A cache layer with signature validation enabled.

    Raises
    ------
    ValueError
        If no secret is provided and TNFR_CACHE_SECRET is not set.

    Examples
    --------
    >>> # set environment variable in production:
    >>> # export TNFR_CACHE_SECRET="your-secure-random-key"
    >>>
    >>> layer = create_secure_redis_layer()
    >>> # Or with explicit configuration:
    >>> import redis
    >>> client = redis.Redis(host='localhost', port=6379)
    >>> layer = create_secure_redis_layer(client, secret=b"my-secret")
    """
    if secret is None:
        secret = os.environ.get("TNFR_CACHE_SECRET")
        if not secret:
            raise TNFRValueError(
                "Secret required for secure cache layer.",
                context={"env_var": "TNFR_CACHE_SECRET"},
                suggestion="set TNFR_CACHE_SECRET environment variable or pass secret parameter.",
            )

    signer = create_hmac_signer(secret)
    validator = create_hmac_validator(secret)

    return RedisCacheLayer(
        client=client,
        namespace=namespace,
        signer=signer,
        validator=validator,
        require_signature=True,
        protocol=protocol,
    )


def _prepare_payload_bytes(value: Any, *, protocol: int) -> tuple[int, bytes]:
    """Return payload encoding mode and the bytes that should be signed."""

    if isinstance(value, (bytes, bytearray, memoryview)):
        return _SIGN_MODE_RAW, bytes(value)
    return _SIGN_MODE_PICKLE, pickle.dumps(value, protocol=protocol)


def _pack_signed_envelope(mode: int, payload: bytes, signature: bytes) -> bytes:
    """Pack payload and signature into a self-describing binary envelope."""

    if not (0 <= mode <= 255):  # pragma: no cover - defensive guard
        raise TNFRValueError(
            f"invalid payload mode: {mode}",
            context={"mode": mode},
            suggestion="Mode must be between 0 and 255.",
        )
    signature_length = len(signature)
    if signature_length >= 2**32:  # pragma: no cover - defensive guard
        raise TNFRValueError(
            "signature too large to encode",
            context={"signature_length": signature_length},
            suggestion="Signature must be smaller than 4GB.",
        )
    header = (
        _SIGNATURE_PREFIX
        + bytes([mode])
        + signature_length.to_bytes(4, byteorder="big", signed=False)
    )
    return header + signature + payload


def _is_signed_envelope(blob: bytes) -> bool:
    """Return ``True`` when *blob* represents a signed cache entry."""

    return blob.startswith((_SIGNATURE_PREFIX, _LEGACY_SIGNATURE_PREFIX))


def _signature_input(mode: int, payload: bytes) -> bytes:
    """Authenticate interpretation as well as payload bytes.

    Version 1 signed only the payload, so a raw byte string could be relabelled
    as a pickle. Those derived cache entries must be rebuilt, not trusted.
    """
    return _SIGNATURE_PREFIX + bytes([mode]) + payload


class _EnvelopeUnpickler(pickle.Unpickler):
    """Defense in depth after primitive-only outer opcode validation."""

    def find_class(self, module: str, name: str) -> Any:
        raise TNFRSecurityError("executable outer shelve pickle rejected")


_OUTER_ENVELOPE_OPCODES = frozenset(
    {
        "PROTO",
        "FRAME",
        "STOP",
        "SHORT_BINBYTES",
        "BINBYTES",
        "BINBYTES8",
        "MEMOIZE",
        "BINPUT",
        "LONG_BINPUT",
        "PUT",
        "BINGET",
        "LONG_BINGET",
        "GET",
    }
)


def _decode_outer_envelope(blob: bytes) -> bytes:
    """Validate a bytes-only outer representation before any unpickling.

    EXT opcodes can bypass Unpickler.find_class through Python's global
    extension cache. No extensions, globals, reducers or persistent IDs are
    permitted in the outer representation, regardless of process state.
    """
    for opcode, _, _ in pickletools.genops(blob):
        if opcode.name not in _OUTER_ENVELOPE_OPCODES:
            raise TNFRSecurityError(f"nonprimitive outer pickle opcode: {opcode.name}")
    entry = _EnvelopeUnpickler(io.BytesIO(blob)).load()
    if type(entry) is not bytes:
        raise TNFRSecurityError("outer shelve envelope must be bytes")
    return entry


def _unpack_signed_envelope(blob: bytes) -> tuple[int, bytes, bytes]:
    """Return the ``(mode, signature, payload)`` triple encoded in *blob*."""

    if blob.startswith(_LEGACY_SIGNATURE_PREFIX):
        raise TNFRSecurityError("legacy signed cache entries must be rebuilt")
    if len(blob) < _SIGNATURE_HEADER_SIZE:
        raise TNFRSecurityError("signed payload header truncated")
    if not _is_signed_envelope(blob):
        raise TNFRSecurityError("missing signed payload marker")
    mode = blob[len(_SIGNATURE_PREFIX)]
    sig_start = len(_SIGNATURE_PREFIX) + 1
    sig_len = int.from_bytes(blob[sig_start : sig_start + 4], byteorder="big")
    payload_start = sig_start + 4 + sig_len
    if len(blob) < payload_start:
        raise TNFRSecurityError("signed payload signature truncated")
    signature = blob[sig_start + 4 : payload_start]
    payload = blob[payload_start:]
    return mode, signature, payload


def _decode_payload(mode: int, payload: bytes) -> Any:
    """Decode payload bytes depending on cache encoding *mode*."""

    if mode == _SIGN_MODE_RAW:
        return payload
    if mode == _SIGN_MODE_PICKLE:
        return pickle.loads(payload)  # nosec B301 - validated via signature
    raise TNFRSecurityError(f"unknown payload encoding mode: {mode}")


class CacheLayer(ABC):
    """Abstract interface implemented by storage backends orchestrated by :class:`CacheManager`."""

    @abstractmethod
    def load(self, name: str) -> Any:
        """Return the stored payload for ``name`` or raise :class:`KeyError`."""

    @abstractmethod
    def store(self, name: str, value: Any) -> None:
        """Persist ``value`` under ``name``."""

    @abstractmethod
    def delete(self, name: str) -> None:
        """Remove ``name`` from the backend if present."""

    @abstractmethod
    def clear(self) -> None:
        """Remove every entry maintained by the layer."""

    def close(self) -> None:  # pragma: no cover - optional hook
        """Release resources held by the backend."""


class MappingCacheLayer(CacheLayer):
    """In-memory cache layer backed by a mutable mapping."""

    def __init__(self, storage: MutableMapping[str, Any] | None = None) -> None:
        self._storage: MutableMapping[str, Any] = {} if storage is None else storage
        self._lock = threading.RLock()

    @property
    def storage(self) -> MutableMapping[str, Any]:
        """Return the mapping used to store cache entries."""

        return self._storage

    def load(self, name: str) -> Any:
        with self._lock:
            if name not in self._storage:
                raise KeyError(name)
            return self._storage[name]

    def store(self, name: str, value: Any) -> None:
        with self._lock:
            self._storage[name] = value

    def delete(self, name: str) -> None:
        with self._lock:
            self._storage.pop(name, None)

    def clear(self) -> None:
        with self._lock:
            self._storage.clear()


class ShelveCacheLayer(CacheLayer):
    """Persistent cache layer backed by :mod:`shelve`.

    .. warning::
        This layer uses :mod:`pickle` for serialization, which can deserialize
        arbitrary Python objects and execute code during deserialization.
        **Only use with trusted data** from controlled sources. Never load
        shelf files from untrusted origins without cryptographic verification.

        Pickle is required for TNFR's complex structures (NetworkX graphs, EPIs,
        coherence states, numpy arrays). For untrusted inputs, enable
        :term:`HMAC` or equivalent signing via ``signer``/``validator`` and
        set ``require_signature=True`` to reject tampered payloads.

    :param signer: Optional callable that receives payload bytes and returns a
        signature (for example ``lambda payload: hmac.new(key, payload,
        hashlib.sha256).digest()``).
    :param validator: Optional callable that receives ``(payload_bytes,
        signature)`` and returns ``True`` when the payload is trustworthy.
    :param require_signature: When ``True`` the cache operates in hardened
        mode, deleting entries whose signatures are missing or invalid and
        raising :class:`SecurityError`.
    """

    def __init__(
        self,
        path: str,
        *,
        flag: str = "c",
        protocol: int | None = None,
        writeback: bool = False,
        signer: Callable[[bytes], bytes] | None = None,
        validator: Callable[[bytes, bytes], bool] | None = None,
        require_signature: bool = False,
    ) -> None:
        # Validate cache file path to prevent path traversal
        from ..security import PathTraversalError, validate_file_path

        try:
            validated_path = validate_file_path(
                path,
                allow_absolute=True,
                allowed_extensions=None,  # Shelve creates multiple files with various extensions
            )
            self._path = str(validated_path)
        except (ValueError, PathTraversalError) as e:
            raise TNFRValueError(
                f"Invalid cache path {path!r}: {e}",
                context={"path": path, "error": str(e)},
            ) from e

        self._flag = flag
        self._protocol = pickle.HIGHEST_PROTOCOL if protocol is None else protocol
        # shelve module inherently uses pickle for serialization; security risks documented in class docstring
        self._shelf = shelve.open(
            self._path, flag=flag, protocol=self._protocol, writeback=writeback  # type: ignore
        )  # nosec B301
        self._lock = threading.RLock()
        self._signer = signer
        self._validator = validator
        self._require_signature = require_signature
        if require_signature and (signer is None or validator is None):
            raise TNFRValueError(
                "require_signature=True requires both signer and validator",
                context={"signer": signer, "validator": validator},
            )

        # Issue security warning when using unsigned pickle deserialization
        if not require_signature and os.environ.get(_TNFR_ALLOW_UNSIGNED_PICKLE) != "1":
            warnings.warn(
                f"ShelveCacheLayer at {path!r} uses pickle without signature validation. "
                "This can execute arbitrary code during deserialization. "
                "Use create_secure_shelve_layer() or set require_signature=True with signer/validator. "
                f"To suppress this warning, set {_TNFR_ALLOW_UNSIGNED_PICKLE}=1 environment variable.",
                TNFRSecurityWarning,
                stacklevel=2,
            )

    def load(self, name: str) -> Any:
        with self._lock:
            if name not in self._shelf:
                raise KeyError(name)
            if self._require_signature or self._validator is not None:
                # Shelf.__getitem__ invokes pickle before returning its value.
                # Verify the outer container cannot execute before examining
                # and authenticating the inner envelope. Unsigned, explicitly
                # trusted shelves retain their original object policy below.
                blob = self._shelf.dict[name.encode(self._shelf.keyencoding)]
                try:
                    entry = _decode_outer_envelope(blob)
                except (
                    pickle.PickleError,
                    EOFError,
                    ValueError,
                    TNFRSecurityError,
                ) as exc:
                    self.delete(name)
                    raise TNFRSecurityError("invalid outer shelve envelope") from exc
            else:
                entry = self._shelf[name]

        return self._decode_entry(name, entry)

    def store(self, name: str, value: Any) -> None:
        if self._signer is None:
            stored_value: Any = value
        else:
            mode, payload = _prepare_payload_bytes(value, protocol=self._protocol)
            signature = self._signer(_signature_input(mode, payload))
            stored_value = _pack_signed_envelope(mode, payload, signature)
        with self._lock:
            if self._signer is None:
                self._shelf[name] = stored_value
            else:
                # A fixed bytes-only outer protocol avoids GLOBAL/REDUCE in
                # protocols 0-2. The requested protocol still controls the
                # authenticated inner value, preserving its public semantics.
                self._shelf.dict[name.encode(self._shelf.keyencoding)] = pickle.dumps(
                    stored_value,
                    protocol=3,
                )
                self._shelf.cache.pop(name, None)
            self._shelf.sync()

    def delete(self, name: str) -> None:
        with self._lock:
            try:
                del self._shelf[name]
            except KeyError:
                return
            self._shelf.sync()

    def clear(self) -> None:
        with self._lock:
            self._shelf.clear()
            self._shelf.sync()

    def close(self) -> None:  # pragma: no cover - exercised indirectly
        with self._lock:
            self._shelf.close()

    def _decode_entry(self, name: str, entry: Any) -> Any:
        if isinstance(entry, (bytes, bytearray, memoryview)):
            blob = bytes(entry)
            if _is_signed_envelope(blob):
                try:
                    mode, signature, payload = _unpack_signed_envelope(blob)
                except TNFRSecurityError:
                    self.delete(name)
                    raise
                validator = self._validator
                if validator is None:
                    if self._require_signature:
                        self.delete(name)
                        raise TNFRSecurityError(
                            "signature validation requested but no validator configured"
                        )
                else:
                    try:
                        valid = validator(_signature_input(mode, payload), signature)
                    except Exception as exc:  # pragma: no cover - defensive
                        self.delete(name)
                        raise TNFRSecurityError(
                            "signature validator raised an exception"
                        ) from exc
                    if not valid:
                        self.delete(name)
                        raise TNFRSecurityError(
                            f"signature validation failed for cache entry {name!r}"
                        )
                try:
                    return _decode_payload(mode, payload)
                except Exception as exc:
                    self.delete(name)
                    raise TNFRSecurityError("signed payload decode failure") from exc
            if self._require_signature:
                self.delete(name)
                raise TNFRSecurityError(f"unsigned cache entry rejected: {name}")
            return blob
        if self._require_signature:
            self.delete(name)
            raise TNFRSecurityError(f"unsigned cache entry rejected: {name}")
        return entry


class RedisCacheLayer(CacheLayer):
    """Distributed cache layer backed by a Redis client.

    .. warning::
        This layer uses :mod:`pickle` for serialization, which can deserialize
        arbitrary Python objects and execute code during deserialization.
        **Only cache trusted data** from controlled TNFR nodes. Ensure Redis
        uses authentication (AUTH command or ACL for Redis 6.0+) and network
        access controls. Never cache untrusted user input or external data.

        If Redis is compromised or contains tampered data, pickle deserialization
        executes arbitrary code. Use TLS for connections and enable signature
        validation (``signer``/``validator`` with ``require_signature=True``)
        in high-assurance deployments.

    :param signer: Optional callable that produces a signature for payload bytes
        before they are written to Redis.
    :param validator: Optional callable that validates ``(payload_bytes,
        signature)`` during loads.
    :param require_signature: Enable hardened mode that deletes and rejects
        cache entries whose signatures are missing or invalid, raising
        :class:`SecurityError`.
    """

    def __init__(
        self,
        client: Any | None = None,
        *,
        namespace: str = "tnfr:cache",
        signer: Callable[[bytes], bytes] | None = None,
        validator: Callable[[bytes, bytes], bool] | None = None,
        require_signature: bool = False,
        protocol: int | None = None,
    ) -> None:
        if client is None:
            try:  # pragma: no cover - import guarded for optional dependency
                import redis  # type: ignore
            except Exception as exc:  # pragma: no cover - defensive import
                raise RuntimeError(
                    "redis-py is required to initialise RedisCacheLayer"
                ) from exc
            client = redis.Redis()
        self._client = client
        self._namespace = namespace.rstrip(":") or "tnfr:cache"
        self._lock = threading.RLock()
        self._signer = signer
        self._validator = validator
        self._require_signature = require_signature
        self._protocol = pickle.HIGHEST_PROTOCOL if protocol is None else protocol
        if require_signature and (signer is None or validator is None):
            raise TNFRValueError(
                "require_signature=True requires both signer and validator",
                context={"signer": signer, "validator": validator},
            )

        # Issue security warning when using unsigned pickle deserialization
        if not require_signature and os.environ.get(_TNFR_ALLOW_UNSIGNED_PICKLE) != "1":
            warnings.warn(
                f"RedisCacheLayer with namespace {namespace!r} uses pickle without signature validation. "
                "This can execute arbitrary code if Redis is compromised. "
                "Use create_secure_redis_layer() or set require_signature=True with signer/validator. "
                f"To suppress this warning, set {_TNFR_ALLOW_UNSIGNED_PICKLE}=1 environment variable.",
                TNFRSecurityWarning,
                stacklevel=2,
            )

    def _format_key(self, name: str) -> str:
        return f"{self._namespace}:{name}"

    def load(self, name: str) -> Any:
        key = self._format_key(name)
        with self._lock:
            value = self._client.get(key)
        if value is None:
            raise KeyError(name)
        if isinstance(value, (bytes, bytearray, memoryview)):
            blob = bytes(value)
            if _is_signed_envelope(blob):
                try:
                    mode, signature, payload = _unpack_signed_envelope(blob)
                except TNFRSecurityError:
                    self.delete(name)
                    raise
                validator = self._validator
                if validator is None:
                    if self._require_signature:
                        self.delete(name)
                        raise TNFRSecurityError(
                            "signature validation requested but no validator configured"
                        )
                else:
                    try:
                        valid = validator(_signature_input(mode, payload), signature)
                    except Exception as exc:  # pragma: no cover - defensive
                        self.delete(name)
                        raise TNFRSecurityError(
                            "signature validator raised an exception"
                        ) from exc
                    if not valid:
                        self.delete(name)
                        raise TNFRSecurityError(
                            f"signature validation failed for cache entry {name!r}"
                        )
                try:
                    return _decode_payload(mode, payload)
                except Exception as exc:
                    self.delete(name)
                    raise TNFRSecurityError("signed payload decode failure") from exc
            if self._require_signature:
                self.delete(name)
                raise TNFRSecurityError(f"unsigned cache entry rejected: {name}")
            # pickle from trusted Redis; documented security warning in class docstring
            return pickle.loads(blob)  # nosec B301
        if self._require_signature:
            self.delete(name)
            raise TNFRSecurityError(f"unsigned nonbinary cache entry rejected: {name}")
        return value

    def store(self, name: str, value: Any) -> None:
        key = self._format_key(name)
        if self._signer is None:
            payload: Any = value
            if not isinstance(value, (bytes, bytearray, memoryview)):
                payload = pickle.dumps(value, protocol=self._protocol)
        else:
            mode, payload_bytes = _prepare_payload_bytes(value, protocol=self._protocol)
            signature = self._signer(_signature_input(mode, payload_bytes))
            payload = _pack_signed_envelope(mode, payload_bytes, signature)
        with self._lock:
            self._client.set(key, payload)

    def delete(self, name: str) -> None:
        key = self._format_key(name)
        with self._lock:
            self._client.delete(key)

    def clear(self) -> None:
        pattern = f"{self._namespace}:*"
        with self._lock:
            if hasattr(self._client, "scan_iter"):
                keys = list(self._client.scan_iter(match=pattern))
            elif hasattr(self._client, "keys"):
                keys = list(self._client.keys(pattern))
            else:  # pragma: no cover - extremely defensive
                keys = []
            if keys:
                self._client.delete(*keys)