merge: testing → main (reconcile 2-week divergence)

decnet/web/db/sqlmodel_repo/canary.py

@@ -0,0 +1,200 @@
+"""Canary blob/token CRUD + trigger ingestion."""
+from __future__ import annotations
+
+import json
+from datetime import datetime, timezone
+from typing import Any, Optional
+
+from sqlalchemy import desc, func, select, update
+
+from decnet.web.db.models import CanaryBlob, CanaryToken, CanaryTrigger
+
+
+class CanaryMixin:
+    """Mixin: composed onto ``SQLModelRepository``."""
+
+    async def upsert_canary_blob(self, data: dict[str, Any]) -> dict[str, Any]:
+        sha = data.get("sha256")
+        if not sha:
+            raise ValueError("upsert_canary_blob: sha256 is required")
+        async with self._session() as session:
+            existing = await session.execute(
+                select(CanaryBlob).where(CanaryBlob.sha256 == sha)
+            )
+            row = existing.scalar_one_or_none()
+            if row:
+                return row.model_dump(mode="json")
+            row = CanaryBlob(**data)
+            session.add(row)
+            await session.commit()
+            await session.refresh(row)
+            return row.model_dump(mode="json")
+
+    async def get_canary_blob(self, uuid: str) -> Optional[dict[str, Any]]:
+        async with self._session() as session:
+            result = await session.execute(
+                select(CanaryBlob).where(CanaryBlob.uuid == uuid)
+            )
+            row = result.scalar_one_or_none()
+            return row.model_dump(mode="json") if row else None
+
+    async def get_canary_blob_by_sha256(
+        self, sha256: str
+    ) -> Optional[dict[str, Any]]:
+        async with self._session() as session:
+            result = await session.execute(
+                select(CanaryBlob).where(CanaryBlob.sha256 == sha256)
+            )
+            row = result.scalar_one_or_none()
+            return row.model_dump(mode="json") if row else None
+
+    async def list_canary_blobs(self) -> list[dict[str, Any]]:
+        # One round-trip: outer-join blobs -> tokens, group by blob, count
+        # live (non-revoked) references.  Revoked tokens still occupy the
+        # blob conceptually until garbage-collected, so we count them too;
+        # the operator deletes blobs explicitly via the API.
+        async with self._session() as session:
+            stmt = (
+                select(CanaryBlob, func.count(CanaryToken.uuid))
+                .join(
+                    CanaryToken,
+                    CanaryToken.blob_uuid == CanaryBlob.uuid,
+                    isouter=True,
+                )
+                .group_by(CanaryBlob.uuid)
+                .order_by(desc(CanaryBlob.uploaded_at))
+            )
+            result = await session.execute(stmt)
+            out: list[dict[str, Any]] = []
+            for blob, count in result.all():
+                d = blob.model_dump(mode="json")
+                d["token_count"] = int(count or 0)
+                out.append(d)
+            return out
+
+    async def delete_canary_blob(self, uuid: str) -> bool:
+        async with self._session() as session:
+            ref = await session.execute(
+                select(func.count(CanaryToken.uuid)).where(
+                    CanaryToken.blob_uuid == uuid
+                )
+            )
+            if (ref.scalar_one() or 0) > 0:
+                return False
+            result = await session.execute(
+                select(CanaryBlob).where(CanaryBlob.uuid == uuid)
+            )
+            row = result.scalar_one_or_none()
+            if not row:
+                return False
+            await session.delete(row)
+            await session.commit()
+            return True
+
+    async def create_canary_token(self, data: dict[str, Any]) -> None:
+        async with self._session() as session:
+            session.add(CanaryToken(**data))
+            await session.commit()
+
+    async def get_canary_token(self, uuid: str) -> Optional[dict[str, Any]]:
+        async with self._session() as session:
+            result = await session.execute(
+                select(CanaryToken).where(CanaryToken.uuid == uuid)
+            )
+            row = result.scalar_one_or_none()
+            return row.model_dump(mode="json") if row else None
+
+    async def get_canary_token_by_slug(
+        self, callback_token: str
+    ) -> Optional[dict[str, Any]]:
+        async with self._session() as session:
+            result = await session.execute(
+                select(CanaryToken).where(
+                    CanaryToken.callback_token == callback_token
+                )
+            )
+            row = result.scalar_one_or_none()
+            return row.model_dump(mode="json") if row else None
+
+    async def list_canary_tokens(
+        self,
+        *,
+        decky_name: Optional[str] = None,
+        state: Optional[str] = None,
+        kind: Optional[str] = None,
+    ) -> list[dict[str, Any]]:
+        async with self._session() as session:
+            stmt = select(CanaryToken)
+            if decky_name is not None:
+                stmt = stmt.where(CanaryToken.decky_name == decky_name)
+            if state is not None:
+                stmt = stmt.where(CanaryToken.state == state)
+            if kind is not None:
+                stmt = stmt.where(CanaryToken.kind == kind)
+            stmt = stmt.order_by(desc(CanaryToken.placed_at))
+            result = await session.execute(stmt)
+            return [r.model_dump(mode="json") for r in result.scalars().all()]
+
+    async def update_canary_token_state(
+        self,
+        uuid: str,
+        state: str,
+        last_error: Optional[str] = None,
+    ) -> bool:
+        async with self._session() as session:
+            result = await session.execute(
+                update(CanaryToken)
+                .where(CanaryToken.uuid == uuid)
+                .values(state=state, last_error=last_error)
+            )
+            await session.commit()
+            return result.rowcount > 0
+
+    async def record_canary_trigger(self, data: dict[str, Any]) -> str:
+        # Persist the trigger row + bump the token's counters in the same
+        # session so a subscriber that reads the token row right after
+        # receiving the bus event sees the updated count.
+        headers = data.get("raw_headers")
+        if isinstance(headers, dict):
+            data = {**data, "raw_headers": json.dumps(headers)}
+        async with self._session() as session:
+            row = CanaryTrigger(**data)
+            session.add(row)
+            ts = data.get("occurred_at") or datetime.now(timezone.utc)
+            await session.execute(
+                update(CanaryToken)
+                .where(CanaryToken.uuid == row.token_uuid)
+                .values(
+                    last_triggered_at=ts,
+                    trigger_count=CanaryToken.trigger_count + 1,
+                )
+            )
+            await session.commit()
+            await session.refresh(row)
+            return row.uuid
+
+    async def list_canary_triggers(
+        self, token_uuid: str, *, limit: int = 100, offset: int = 0,
+    ) -> list[dict[str, Any]]:
+        async with self._session() as session:
+            stmt = (
+                select(CanaryTrigger)
+                .where(CanaryTrigger.token_uuid == token_uuid)
+                .order_by(desc(CanaryTrigger.occurred_at))
+                .limit(limit)
+                .offset(offset)
+            )
+            result = await session.execute(stmt)
+            return [r.model_dump(mode="json") for r in result.scalars().all()]
+
+    async def attribute_canary_trigger(
+        self, trigger_uuid: str, attacker_id: str,
+    ) -> bool:
+        async with self._session() as session:
+            result = await session.execute(
+                update(CanaryTrigger)
+                .where(CanaryTrigger.uuid == trigger_uuid)
+                .values(attacker_id=attacker_id)
+            )
+            await session.commit()
+            return result.rowcount > 0