import asyncio import logging from abc import ABC, abstractmethod from enum import StrEnum logger = logging.getLogger(__name__) class AllowlistHitReason(StrEnum): DM_WILDCARD = "dm_wildcard" DM_EXACT = "dm_exact" GROUP_WILDCARD = "group_wildcard" GROUP_EXACT = "group_exact" GROUP_ALLOW_FROM_EXACT = "group_allow_from_exact" REJECTED = "rejected" DM_OPEN = "dm_open" DM_DISABLED = "dm_disabled" GROUP_OPEN = "group_open" GROUP_DISABLED = "group_disabled" ALLOW_FROM_EMPTY = "allow_from_empty" class DmPolicyMode(StrEnum): OPEN = "open" PAIRING = "pairing" ALLOWLIST = "allowlist" DISABLED = "disabled" class GroupPolicyMode(StrEnum): OPEN = "open" ALLOWLIST = "allowlist" DISABLED = "disabled" class AllowlistCheckResult: def __init__( self, allowed: bool, reason: AllowlistHitReason = AllowlistHitReason.REJECTED, matched_entry: str | None = None, pairing_required: bool = False, ): self.allowed = allowed self.reason = reason self.matched_entry = matched_entry self.pairing_required = pairing_required # ── AllowlistStore (abstract) ───────────────────────────────────────── class AllowlistStore(ABC): @abstractmethod async def load_all(self, channel_id: str) -> dict[str, list[str]]: ... @abstractmethod async def add_entry(self, channel_id: str, allow_type: str, allow_id: str) -> None: ... @abstractmethod async def remove_entry(self, channel_id: str, allow_type: str, allow_id: str) -> None: ... # ── InMemoryAllowlistStore ──────────────────────────────────────────── class InMemoryAllowlistStore(AllowlistStore): def __init__(self): self._entries: dict[str, dict[str, set[str]]] = {} async def load_all(self, channel_id: str) -> dict[str, list[str]]: channel_entries = self._entries.get(channel_id, {}) return {allow_type: list(ids) for allow_type, ids in channel_entries.items()} async def add_entry(self, channel_id: str, allow_type: str, allow_id: str) -> None: if channel_id not in self._entries: self._entries[channel_id] = {} if allow_type not in self._entries[channel_id]: self._entries[channel_id][allow_type] = set() self._entries[channel_id][allow_type].add(allow_id) async def remove_entry(self, channel_id: str, allow_type: str, allow_id: str) -> None: if channel_id in self._entries and allow_type in self._entries[channel_id]: self._entries[channel_id][allow_type].discard(allow_id) # ── PostgresAllowlistStore ──────────────────────────────────────────── class PostgresAllowlistStore(AllowlistStore): def __init__(self): from yuxi.repositories.channel_allowlist_repo import ChannelAllowlistRepository self._repo = ChannelAllowlistRepository() async def load_all(self, channel_id: str) -> dict[str, list[str]]: entries = await self._repo.list_by_channel(channel_id) result: dict[str, list[str]] = {} for entry in entries: result.setdefault(entry.allow_type, []).append(entry.allow_id) return result async def add_entry(self, channel_id: str, allow_type: str, allow_id: str) -> None: await self._repo.add_entry(channel_id, allow_type, allow_id) async def remove_entry(self, channel_id: str, allow_type: str, allow_id: str) -> None: await self._repo.remove_entry(channel_id, allow_type, allow_id) # ── AllowlistChecker ────────────────────────────────────────────────── class AllowlistChecker: def __init__(self, store: AllowlistStore | None = None): self._store = store or InMemoryAllowlistStore() self._dm_allowlist: list[str] = [] self._group_allowlist: list[str] = [] self._group_allow_from: list[str] = [] self._dm_policy: DmPolicyMode = DmPolicyMode.ALLOWLIST self._group_policy: GroupPolicyMode = GroupPolicyMode.ALLOWLIST self._loaded = False self._loading = False self._channel_id: str | None = None self._load_lock = asyncio.Lock() def load_dm_allowlist(self, entries: list[str]) -> None: self._dm_allowlist = list(entries) self._loaded = True def load_group_allowlist(self, group_ids: list[str]) -> None: self._group_allowlist = list(group_ids) self._loaded = True def load_group_allow_from(self, entries: list[str]) -> None: self._group_allow_from = list(entries) async def ensure_loaded(self, channel_id: str) -> None: if self._loaded: return async with self._load_lock: if self._loaded: return self._loading = True try: data = await self._store.load_all(channel_id) self._dm_allowlist = data.get("user", []) self._group_allowlist = data.get("group", []) self._group_allow_from = data.get("allow_from", []) self._loaded = True self._channel_id = channel_id except Exception: logger.exception("Failed to load allowlist for channel %s", channel_id) raise finally: self._loading = False @property def dm_policy(self) -> DmPolicyMode: return self._dm_policy @dm_policy.setter def dm_policy(self, mode: DmPolicyMode) -> None: self._dm_policy = mode @property def group_policy(self) -> GroupPolicyMode: return self._group_policy @group_policy.setter def group_policy(self, mode: GroupPolicyMode) -> None: self._group_policy = mode def check_dm(self, peer_id: str) -> AllowlistCheckResult: if self._dm_policy == DmPolicyMode.OPEN: return AllowlistCheckResult(allowed=True, reason=AllowlistHitReason.DM_OPEN) if self._dm_policy == DmPolicyMode.DISABLED: return AllowlistCheckResult(allowed=False, reason=AllowlistHitReason.DM_DISABLED) if self._dm_policy == DmPolicyMode.PAIRING: if self._dm_allowlist and "*" in self._dm_allowlist: return AllowlistCheckResult(allowed=True, reason=AllowlistHitReason.DM_WILDCARD, matched_entry="*") if peer_id in self._dm_allowlist: return AllowlistCheckResult(allowed=True, reason=AllowlistHitReason.DM_EXACT, matched_entry=peer_id) return AllowlistCheckResult( allowed=False, reason=AllowlistHitReason.REJECTED, pairing_required=True, ) if not self._dm_allowlist: return AllowlistCheckResult(allowed=False, reason=AllowlistHitReason.REJECTED) if "*" in self._dm_allowlist: return AllowlistCheckResult(allowed=True, reason=AllowlistHitReason.DM_WILDCARD, matched_entry="*") if peer_id in self._dm_allowlist: return AllowlistCheckResult(allowed=True, reason=AllowlistHitReason.DM_EXACT, matched_entry=peer_id) return AllowlistCheckResult(allowed=False, reason=AllowlistHitReason.REJECTED) def check_group(self, group_id: str) -> AllowlistCheckResult: if self._group_policy == GroupPolicyMode.OPEN: return AllowlistCheckResult(allowed=True, reason=AllowlistHitReason.GROUP_OPEN) if self._group_policy == GroupPolicyMode.DISABLED: return AllowlistCheckResult(allowed=False, reason=AllowlistHitReason.GROUP_DISABLED) if not self._group_allowlist: return AllowlistCheckResult(allowed=False, reason=AllowlistHitReason.REJECTED) if "*" in self._group_allowlist: return AllowlistCheckResult(allowed=True, reason=AllowlistHitReason.GROUP_WILDCARD, matched_entry="*") if group_id in self._group_allowlist: return AllowlistCheckResult(allowed=True, reason=AllowlistHitReason.GROUP_EXACT, matched_entry=group_id) return AllowlistCheckResult(allowed=False, reason=AllowlistHitReason.REJECTED) def check_group_allow_from(self, sender_id: str) -> AllowlistCheckResult: if self._group_policy == GroupPolicyMode.OPEN: return AllowlistCheckResult(allowed=True, reason=AllowlistHitReason.GROUP_OPEN) if self._group_policy == GroupPolicyMode.DISABLED: return AllowlistCheckResult(allowed=False, reason=AllowlistHitReason.GROUP_DISABLED) if not self._group_allow_from: logger.warning( "group_allow_from is empty under allowlist policy for channel %s — blocking all group senders. " "Add explicit entries or use groupPolicy=open to allow.", self._channel_id, ) return AllowlistCheckResult(allowed=False, reason=AllowlistHitReason.ALLOW_FROM_EMPTY) if "*" in self._group_allow_from: return AllowlistCheckResult( allowed=True, reason=AllowlistHitReason.GROUP_ALLOW_FROM_EXACT, matched_entry="*", ) if sender_id in self._group_allow_from: return AllowlistCheckResult( allowed=True, reason=AllowlistHitReason.GROUP_ALLOW_FROM_EXACT, matched_entry=sender_id, ) return AllowlistCheckResult(allowed=False, reason=AllowlistHitReason.REJECTED) def add_to_dm_allowlist(self, peer_id: str) -> None: if peer_id not in self._dm_allowlist: self._dm_allowlist.append(peer_id) def remove_from_dm_allowlist(self, peer_id: str) -> None: if peer_id in self._dm_allowlist: self._dm_allowlist.remove(peer_id) _VALID_SCOPES = frozenset({"dm", "group", "allow_from"}) async def add_entry(self, channel_type: str, scope: str, value: str) -> bool: if scope not in self._VALID_SCOPES: raise ValueError(f"Invalid scope: {scope!r}, must be one of {sorted(self._VALID_SCOPES)}") allow_type = {"dm": "user", "group": "group", "allow_from": "allow_from"}[scope] await self._store.add_entry(channel_type, allow_type, value) if scope == "dm": if value not in self._dm_allowlist: self._dm_allowlist.append(value) return True if scope == "group": if value not in self._group_allowlist: self._group_allowlist.append(value) return True if scope == "allow_from": if value not in self._group_allow_from: self._group_allow_from.append(value) return True return False async def remove_entry(self, channel_type: str, scope: str, value: str) -> bool: if scope not in self._VALID_SCOPES: raise ValueError(f"Invalid scope: {scope!r}, must be one of {sorted(self._VALID_SCOPES)}") allow_type = {"dm": "user", "group": "group", "allow_from": "allow_from"}[scope] await self._store.remove_entry(channel_type, allow_type, value) if scope == "dm": if value in self._dm_allowlist: self._dm_allowlist.remove(value) return True if scope == "group": if value in self._group_allowlist: self._group_allowlist.remove(value) return True if scope == "allow_from": if value in self._group_allow_from: self._group_allow_from.remove(value) return True return False def list_entries(self, channel_type: str, scope: str) -> list[str]: if scope == "dm": return list(self._dm_allowlist) if scope == "group": return list(self._group_allowlist) if scope == "allow_from": return list(self._group_allow_from) return [] # ── Singleton ───────────────────────────────────────────────────────── import contextvars _allowlist_ctx: contextvars.ContextVar[AllowlistChecker | None] = contextvars.ContextVar( "allowlist_checker", default=None ) _allowlist_lock = asyncio.Lock() def create_allowlist_checker() -> AllowlistChecker: return AllowlistChecker(PostgresAllowlistStore()) async def get_allowlist_checker() -> AllowlistChecker: checker = _allowlist_ctx.get() if checker is None: async with _allowlist_lock: if _allowlist_ctx.get() is None: _allowlist_ctx.set(create_allowlist_checker()) checker = _allowlist_ctx.get() if checker is None: checker = create_allowlist_checker() _allowlist_ctx.set(checker) return checker def set_allowlist_checker(checker: AllowlistChecker) -> None: _allowlist_ctx.set(checker) async def reset_allowlist_checker() -> None: async with _allowlist_lock: _allowlist_ctx.set(None)