memory/campaign_store.py

from __future__ import annotations

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

from shield.memory.entity_store import EncryptedEntityStore


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


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

    def get(self, campaign_id: str) -> Dict[str, Any]:
        return self.store.read("campaigns", f"{campaign_id}.json") or {}

    def update(
        self,
        campaign_id: str,
        *,
        account_hash: str,
        entity_hashes: Dict[str, str],
        risk_score: int,
    ) -> Dict[str, Any]:
        record = self.get(campaign_id) or {
            "campaign_id": campaign_id,
            "first_seen": _now(),
            "last_seen": None,
            "account_hashes": [],
            "device_hashes": [],
            "ip_hashes": [],
            "risk_score": 0,
            "seen_count": 0,
        }
        record["last_seen"] = _now()
        record["seen_count"] = int(record.get("seen_count", 0)) + 1
        if account_hash not in record["account_hashes"]:
            record["account_hashes"].append(account_hash)
        for entity_type, bucket in (("device", "device_hashes"), ("ip", "ip_hashes")):
            value = entity_hashes.get(entity_type)
            if value and value not in record[bucket]:
                record[bucket].append(value)
        record["account_count"] = len(record["account_hashes"])
        record["device_count"] = len(record["device_hashes"])
        record["ip_count"] = len(record["ip_hashes"])
        record["risk_score"] = max(int(record.get("risk_score", 0)), int(risk_score))
        self.store.write(record, "campaigns", f"{campaign_id}.json")
        return record