memory/global_memory.py

from __future__ import annotations

from datetime import datetime, timezone
from typing import Any, Dict

from shield.memory.entity_store import EncryptedEntityStore
from shield.utils.hashing import hash_entity


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


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

    def lookup(self, entities: Dict[str, str]) -> Dict[str, Dict[str, Any]]:
        results: Dict[str, Dict[str, Any]] = {}
        for entity_type, value in entities.items():
            entity_hash = hash_entity(entity_type, value)
            record = self.store.read("global", entity_type, f"{entity_hash}.json")
            if record:
                results[entity_type] = record
        return results

    def update(
        self,
        entities: Dict[str, str],
        *,
        customer_id: str | None,
        risk_score: int,
        decision: str,
    ) -> None:
        customer_hash = hash_entity("customer", customer_id or "anonymous")
        for entity_type, value in entities.items():
            entity_hash = hash_entity(entity_type, value)
            record = self.store.read("global", entity_type, f"{entity_hash}.json") or {
                "entity_hash": entity_hash,
                "entity_type": entity_type,
                "first_seen": _now(),
                "last_seen": None,
                "seen_count": 0,
                "abuse_count": 0,
                "avg_risk": 0,
                "linked_customer_hashes": [],
                "flags": [],
            }
            seen_count = int(record.get("seen_count", 0)) + 1
            previous_avg = float(record.get("avg_risk", 0))
            record["seen_count"] = seen_count
            record["avg_risk"] = round(((previous_avg * (seen_count - 1)) + risk_score) / seen_count, 2)
            record["last_seen"] = _now()
            if decision in {"block", "review", "rate_limit"}:
                record["abuse_count"] = int(record.get("abuse_count", 0)) + 1
                if decision not in record["flags"]:
                    record["flags"].append(decision)
            if customer_hash not in record["linked_customer_hashes"]:
                record["linked_customer_hashes"].append(customer_hash)
            record["linked_customers"] = len(record["linked_customer_hashes"])
            self.store.write(record, "global", entity_type, f"{entity_hash}.json")