memory/customer_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 CustomerMemory:
    def __init__(self, store: EncryptedEntityStore | None = None) -> None:
        self.store = store or EncryptedEntityStore()

    def _customer_path(self, customer_id: str | None) -> str:
        return hash_entity("customer", customer_id or "anonymous")

    def lookup(self, customer_id: str | None) -> Dict[str, Any]:
        return self.store.read("customers", self._customer_path(customer_id), "memory.json") or {}

    def update(
        self,
        customer_id: str | None,
        entities: Dict[str, str],
        *,
        account_hash: str,
        risk_score: int,
        decision: str,
        reasons: list[str],
    ) -> None:
        path_id = self._customer_path(customer_id)
        record = self.lookup(customer_id) or {
            "known_devices": [],
            "known_ips": [],
            "known_emails": [],
            "known_phone_numbers": [],
            "known_usernames": [],
            "abuse_events": [],
            "linked_accounts": [],
        }
        mapping = {
            "device": "known_devices",
            "ip": "known_ips",
            "email": "known_emails",
            "phone": "known_phone_numbers",
            "username": "known_usernames",
        }
        for entity_type, value in entities.items():
            bucket = mapping.get(entity_type)
            if not bucket:
                continue
            entity_hash = hash_entity(entity_type, value)
            if entity_hash not in record[bucket]:
                record[bucket].append(entity_hash)
        if account_hash not in record["linked_accounts"]:
            record["linked_accounts"].append(account_hash)
        if decision in {"block", "review", "rate_limit"}:
            record["abuse_events"].append({
                "at": _now(),
                "account_hash": account_hash,
                "risk_score": risk_score,
                "decision": decision,
                "reasons": reasons[:5],
            })
            record["abuse_events"] = record["abuse_events"][-100:]
        self.store.write(record, "customers", path_id, "memory.json")