memory/entity_store.py

from __future__ import annotations

from typing import Any, Dict, Optional

from shield.memory.encryption import decrypt_json, encrypt_json
from shield.utils.storage import JsonStorage


class EncryptedEntityStore:
    def __init__(self, storage: JsonStorage | None = None) -> None:
        self.storage = storage or JsonStorage()

    def read(self, *parts: str) -> Optional[Dict[str, Any]]:
        envelope = self.storage.read_json(*parts)
        if not envelope:
            return None
        return decrypt_json(envelope)

    def write(self, data: Dict[str, Any], *parts: str) -> bool:
        envelope = encrypt_json(data)
        if envelope is None:
            return False
        return self.storage.write_json(envelope, *parts)