diff --git a/.env.example b/.env.example index 2747b7a..6e41ac3 100644 --- a/.env.example +++ b/.env.example @@ -78,6 +78,10 @@ TASK_WORKER_ENABLED=true # Disabled in the full sample. Enable only after adding a real token. TELEGRAM_ENABLED=false # TELEGRAM_BOT_TOKEN=replace-me +# Optional shared gate for new private chats. Existing private chats with +# persisted Blacki history are grandfathered; groups are always rejected. +# New users authenticate with: /start +# TELEGRAM_ACCESS_CODE=replace-with-a-dedicated-high-entropy-code # Telegram turns automatically show one live-updating tool status message. # Optional private text-to-speech delivery for Telegram replies. The Kokoro diff --git a/docs/base-infra/environment-variables.md b/docs/base-infra/environment-variables.md index 62c0a43..965de62 100644 --- a/docs/base-infra/environment-variables.md +++ b/docs/base-infra/environment-variables.md @@ -93,10 +93,15 @@ its presence changes model routing. | --- | --- | --- | | `TELEGRAM_ENABLED` | `false` | Start Telegram long polling | | `TELEGRAM_BOT_TOKEN` | unset | Token from BotFather | +| `TELEGRAM_ACCESS_CODE` | unset | Shared code required for new private Telegram chats | | `KOKORO_TTS_BASE_URL` | unset | Register private Kokoro speech delivery for Telegram | | `KOKORO_TTS_VOICE` | `af_heart` | Kokoro voice ID used for generated MP3 audio | -The token is required and format-validated when Telegram is enabled. +The token is required and format-validated when Telegram is enabled. When +`TELEGRAM_ACCESS_CODE` is set, new users enter it with `/start `; +historical private chats with existing Blacki sessions remain authorized, while +groups and topics are rejected. Rotating the code requires code-authorized +users to authenticate again without deleting their stored Blacki data. `KOKORO_TTS_BASE_URL` is an optional HTTP or HTTPS base URL without a path to `/v1/audio/speech`; Blacki appends that fixed endpoint. The URL must be diff --git a/docs/telegram-setup.md b/docs/telegram-setup.md index 70cde67..0ec0d40 100644 --- a/docs/telegram-setup.md +++ b/docs/telegram-setup.md @@ -22,11 +22,22 @@ In `.env`: ```dotenv TELEGRAM_ENABLED=true TELEGRAM_BOT_TOKEN=replace-me +# Optional: restrict new private chats with one shared access code. +# TELEGRAM_ACCESS_CODE=replace-with-a-dedicated-high-entropy-code ``` Replace `replace-me` with the token from BotFather. Blacki validates that the token is present and follows Telegram's `number:string` format at startup. +When `TELEGRAM_ACCESS_CODE` is set, new users must send `/start ` +in a private chat before Blacki processes their messages. The bot consumes this +command locally and attempts to delete it, but Telegram may retain it in +server-side history or backups, so use a dedicated code rather than a password +you use elsewhere. Historical private chats with persisted Blacki sessions are +grandfathered; group chats and forum topics are rejected. Changing the access +code requires passphrase-authorized users to authenticate again but does not +delete any chat history, preferences, reminders, files, or health data. + At least one model provider must also be configured. See [Configuration](base-infra/environment-variables.md). diff --git a/src/blacki/adk_runtime.py b/src/blacki/adk_runtime.py index ab26b21..4889dd8 100644 --- a/src/blacki/adk_runtime.py +++ b/src/blacki/adk_runtime.py @@ -360,6 +360,10 @@ async def create_next_session( state=state, ) + async def has_existing_session(self, *, locator: SessionLocator) -> bool: + """Return whether a locator has a persisted session without creating one.""" + return await self._get_latest_session(locator=locator) is not None + async def run_user_turn( self, *, diff --git a/src/blacki/container.py b/src/blacki/container.py index 9f3cea1..fe17ece 100644 --- a/src/blacki/container.py +++ b/src/blacki/container.py @@ -30,6 +30,7 @@ from blacki.declarative_db.storage import SqliteDeclarativeDbStorage from blacki.health.storage import SqliteGoogleHealthStorage from blacki.reminders.storage import SqliteReminderStorage + from blacki.telegram.access import TelegramAccessStorage from blacki.user_files.storage import SqliteUserFileStorage from blacki.utils.preferences import SqlitePreferencesStorage from blacki.workouts.storage import SqliteWorkoutStorage @@ -146,6 +147,9 @@ class AppContainer: _user_file_storage: SqliteUserFileStorage | None = field( default=None, init=False, repr=False ) + _telegram_access_storage: TelegramAccessStorage | None = field( + default=None, init=False, repr=False + ) @classmethod async def create(cls, sqlite_path: str | Path) -> Self: @@ -198,6 +202,10 @@ async def _close_storages(self) -> None: await self._user_file_storage.close() self._user_file_storage = None + if self._telegram_access_storage is not None: + await self._telegram_access_storage.close() + self._telegram_access_storage = None + async def initialize_all_storages(self) -> None: """Initialize all storage instances. @@ -211,6 +219,7 @@ async def initialize_all_storages(self) -> None: await self.declarative_db_storage.initialize() await self.google_health_storage.initialize() await self.user_file_storage.initialize() + await self.telegram_access_storage.initialize() @property def lock(self) -> asyncio.Lock: @@ -283,3 +292,12 @@ def user_file_storage(self) -> SqliteUserFileStorage: self._user_file_storage = SqliteUserFileStorage(self.conn, self._lock) return self._user_file_storage + + @property + def telegram_access_storage(self) -> TelegramAccessStorage: + """Get or create local Telegram access and identity storage.""" + if self._telegram_access_storage is None: + from blacki.telegram.access import TelegramAccessStorage + + self._telegram_access_storage = TelegramAccessStorage(self.conn, self._lock) + return self._telegram_access_storage diff --git a/src/blacki/dashboard/data.py b/src/blacki/dashboard/data.py index 16dc8db..99a4f11 100644 --- a/src/blacki/dashboard/data.py +++ b/src/blacki/dashboard/data.py @@ -42,6 +42,7 @@ ) _SESSION_VERSION_RE = re.compile(r"^(?P.+)-v(?P[0-9]+)$") +_TELEGRAM_DIRECT_USER_ID_RE = re.compile(r"^telegram-chat-(?P[0-9]+)$") _ABSOLUTE_PATH_RE = re.compile( r"(? _Snapshot: ) +def _load_telegram_identities_sync(path: Path) -> dict[int, tuple[str, str | None]]: + """Read dashboard-only Telegram labels without modifying the tools database.""" + if not path.is_file(): + return {} + connection: sqlite3.Connection | None = None + try: + uri = f"file:{quote(path.resolve().as_posix(), safe='/')}?mode=ro" + connection = sqlite3.connect(uri, uri=True, timeout=1.0, isolation_level=None) + connection.row_factory = sqlite3.Row + connection.execute("PRAGMA query_only = ON") + rows = connection.execute( + """ + SELECT telegram_user_id, display_name, username + FROM telegram_identities + """ + ).fetchall() + return { + int(row["telegram_user_id"]): (str(row["display_name"]), row["username"]) + for row in rows + if isinstance(row["telegram_user_id"], int) + and isinstance(row["display_name"], str) + and (row["username"] is None or isinstance(row["username"], str)) + } + except (OSError, sqlite3.Error, ValueError): + return {} + finally: + if connection is not None: + connection.close() + + def _iter_raw_jsonl_lines(handle: Any) -> Iterator[tuple[bytes, bool]]: """Yield bounded lines and flag oversized lines without allocating them.""" buffer = bytearray() @@ -1166,10 +1197,12 @@ def __init__( session_db_path: Path, log_dir: Path, app_name: str = "blacki", + identity_db_path: Path | None = None, ) -> None: self.session_db_path = Path(session_db_path) self.log_dir = Path(log_dir) self.app_name = str(app_name) + self.identity_db_path = Path(identity_db_path) if identity_db_path else None async def _snapshot(self) -> _Snapshot: return await asyncio.to_thread( @@ -1321,6 +1354,13 @@ async def get_overview(self, window: str) -> JsonObject: async def list_users(self, search: str, limit: int, offset: int) -> JsonObject: snapshot = await self._snapshot() + identities = ( + await asyncio.to_thread( + _load_telegram_identities_sync, self.identity_db_path + ) + if self.identity_db_path is not None + else {} + ) query = bounded_search(search).lower() page_limit = clamp_limit(limit) page_offset = clamp_offset(offset) @@ -1330,7 +1370,22 @@ async def list_users(self, search: str, limit: int, offset: int) -> JsonObject: ) items: list[JsonObject] = [] for user_id in user_ids: - if query and query not in user_id.lower(): + identity_match = _TELEGRAM_DIRECT_USER_ID_RE.match(user_id) + identity = ( + identities.get(int(identity_match.group("user_id"))) + if identity_match is not None + else None + ) + display_name = identity[0] if identity else None + username = identity[1] if identity else None + searchable_identity = " ".join( + item for item in (display_name, username) if item + ).lower() + if ( + query + and query not in user_id.lower() + and query not in searchable_identity + ): continue sessions = [ session for session in snapshot.sessions if session.user_id == user_id @@ -1349,6 +1404,8 @@ async def list_users(self, search: str, limit: int, offset: int) -> JsonObject: items.append( { "user_id": user_id, + "display_name": display_name, + "username": username, "session_count": len(sessions), "reset_count": max(0, max(versions, default=1) - 1), "retained_history": len(sessions) > 1, diff --git a/src/blacki/dashboard/routes.py b/src/blacki/dashboard/routes.py index 4623245..3fb64af 100644 --- a/src/blacki/dashboard/routes.py +++ b/src/blacki/dashboard/routes.py @@ -217,6 +217,15 @@ def _session_db_path(env: ServerEnv) -> Path: return Path(env.agent_dir) / ".adk" / "sessions.db" +def _tools_db_path(env: ServerEnv) -> Path: + """Return the local SQLite database containing dashboard identity labels.""" + return ( + Path(env.sqlite_path) + if env.sqlite_path + else Path(env.agent_dir) / ".adk" / "tools.db" + ) + + def create_dashboard_router( env: ServerEnv, store: DashboardStoreProtocol | None = None, @@ -235,6 +244,7 @@ def create_dashboard_router( _session_db_path(env), get_log_dir(), "blacki", + _tools_db_path(env), ) except Exception: logger.exception("Dashboard store initialization failed") diff --git a/src/blacki/dashboard/static/dashboard.js b/src/blacki/dashboard/static/dashboard.js index 812c948..13b2170 100644 --- a/src/blacki/dashboard/static/dashboard.js +++ b/src/blacki/dashboard/static/dashboard.js @@ -294,11 +294,16 @@ target.replaceChildren(); state.users.forEach((item) => { const userId = String(first(item, ["user_id", "userId", "chat_id", "chatId", "id"], "—")); + const displayName = first(item, ["display_name", "displayName"]); + const username = first(item, ["username"]); const button = el("button", "btn btn-ghost h-auto min-h-0 w-full justify-between gap-3 rounded-xl px-3 py-3 text-left normal-case hover:bg-base-200"); button.type = "button"; button.setAttribute("aria-label", `Inspect user ${userId}`); const copy = el("span", "min-w-0"); - copy.append(el("span", "block break-all font-mono text-xs font-semibold", userId)); + copy.append(el("span", "block break-all font-mono text-xs font-semibold", displayName || userId)); + if (displayName) { + copy.append(el("span", "mt-1 block break-all font-mono text-xs text-base-content/55", username ? `${userId} · @${username}` : userId)); + } const lastSeen = first(item, ["last_seen", "lastSeen", "updated_at", "updatedAt", "latest_update_at"]); copy.append(el("span", "mt-1 block text-xs text-base-content/55", lastSeen === undefined ? "Stored user ID" : `Seen ${formatDate(lastSeen)}`)); const count = first(item, ["sessions", "session_count", "message_count", "messages"]); diff --git a/src/blacki/server.py b/src/blacki/server.py index f19f6f7..0c4c7c8 100644 --- a/src/blacki/server.py +++ b/src/blacki/server.py @@ -107,6 +107,7 @@ async def _start_telegram_bot() -> None: { "TELEGRAM_ENABLED": env.telegram_enabled, "TELEGRAM_BOT_TOKEN": env.telegram_bot_token, + "TELEGRAM_ACCESS_CODE": env.telegram_access_code, } ) telegram_app = create_app( diff --git a/src/blacki/telegram/__init__.py b/src/blacki/telegram/__init__.py index 2f23395..9799846 100644 --- a/src/blacki/telegram/__init__.py +++ b/src/blacki/telegram/__init__.py @@ -29,6 +29,12 @@ class TelegramConfig(BaseModel): description="Telegram bot token obtained from @BotFather", ) + telegram_access_code: str | None = Field( + default=None, + alias="TELEGRAM_ACCESS_CODE", + description="Shared access code required for new private Telegram users", + ) + model_config = ConfigDict( populate_by_name=True, extra="ignore", @@ -41,3 +47,8 @@ def is_configured(self) -> bool: True if enabled and has a bot token, False otherwise. """ return self.telegram_enabled and self.telegram_bot_token is not None + + @property + def access_control_enabled(self) -> bool: + """Return whether new-user access control is configured.""" + return bool(self.telegram_access_code) diff --git a/src/blacki/telegram/access.py b/src/blacki/telegram/access.py new file mode 100644 index 0000000..e1bf14e --- /dev/null +++ b/src/blacki/telegram/access.py @@ -0,0 +1,130 @@ +"""Local Telegram access control and identity storage.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Literal + +from blacki.storage.base import SqlStorage +from blacki.utils.timezone import now_utc + +if TYPE_CHECKING: + import asyncio + + import aiosqlite + + +AuthorizationSource = Literal["legacy", "passphrase"] + + +@dataclass(frozen=True, slots=True) +class TelegramIdentity: + """A locally stored, user-controlled Telegram display identity.""" + + user_id: int + display_name: str + username: str | None + + +class TelegramAccessStorage(SqlStorage): + """Persist Telegram authorization and dashboard-only identity labels.""" + + def __init__(self, conn: aiosqlite.Connection, lock: asyncio.Lock) -> None: + super().__init__(conn, lock) + + async def _create_tables(self) -> None: + await self._conn.executescript(""" + CREATE TABLE IF NOT EXISTS telegram_access ( + telegram_user_id INTEGER PRIMARY KEY, + source TEXT NOT NULL CHECK(source IN ('legacy', 'passphrase')), + access_code_fingerprint TEXT, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL + ); + CREATE TABLE IF NOT EXISTS telegram_identities ( + telegram_user_id INTEGER PRIMARY KEY, + display_name TEXT NOT NULL, + username TEXT, + updated_at TEXT NOT NULL + ); + """) + + async def is_authorized( + self, telegram_user_id: int, access_code_fingerprint: str + ) -> bool: + row = await self._fetch_one( + """ + SELECT source, access_code_fingerprint + FROM telegram_access + WHERE telegram_user_id = ? + """, + (telegram_user_id,), + ) + if row is None: + return False + source = row["source"] + fingerprint = row["access_code_fingerprint"] + return source == "legacy" or ( + source == "passphrase" + and isinstance(fingerprint, str) + and fingerprint == access_code_fingerprint + ) + + async def has_authorization_record(self, telegram_user_id: int) -> bool: + """Return whether this user was previously granted any authorization.""" + return ( + await self._fetch_one( + "SELECT 1 FROM telegram_access WHERE telegram_user_id = ?", + (telegram_user_id,), + ) + is not None + ) + + async def grant( + self, + telegram_user_id: int, + *, + source: AuthorizationSource, + access_code_fingerprint: str | None = None, + ) -> None: + now = now_utc().isoformat(timespec="seconds") + async with self._lock: + await self._conn.execute( + """ + INSERT INTO telegram_access ( + telegram_user_id, source, access_code_fingerprint, + created_at, updated_at + ) VALUES (?, ?, ?, ?, ?) + ON CONFLICT(telegram_user_id) DO UPDATE SET + source = excluded.source, + access_code_fingerprint = excluded.access_code_fingerprint, + updated_at = excluded.updated_at + """, + (telegram_user_id, source, access_code_fingerprint, now, now), + ) + + async def record_identity(self, identity: TelegramIdentity) -> None: + now = now_utc().isoformat(timespec="seconds") + async with self._lock: + await self._conn.execute( + """ + INSERT INTO telegram_identities ( + telegram_user_id, display_name, username, updated_at + ) VALUES (?, ?, ?, ?) + ON CONFLICT(telegram_user_id) DO UPDATE SET + display_name = excluded.display_name, + username = excluded.username, + updated_at = excluded.updated_at + """, + (identity.user_id, identity.display_name, identity.username, now), + ) + + +def get_telegram_access_storage() -> TelegramAccessStorage: + """Return the initialized process-wide Telegram access storage.""" + from blacki.container import get_container + + storage = get_container().telegram_access_storage + if not storage.is_initialized: + raise RuntimeError("Telegram access storage is not initialized") + return storage diff --git a/src/blacki/telegram/bot.py b/src/blacki/telegram/bot.py index fee2269..3171f79 100644 --- a/src/blacki/telegram/bot.py +++ b/src/blacki/telegram/bot.py @@ -2,6 +2,8 @@ import asyncio import contextlib +import hashlib +import hmac import logging import re from collections.abc import Coroutine, Sequence @@ -28,6 +30,7 @@ from blacki.utils.preferences import get_preferences_storage from . import TelegramConfig +from .access import TelegramAccessStorage, TelegramIdentity, get_telegram_access_storage from .album_buffer import AlbumBuffer, _BufferedAlbum from .api import TelegramApiClient, TelegramApiError from .formatting import escape_markdown_plain, format_for_telegram @@ -42,6 +45,7 @@ Message, ParseMode, Update, + User, ) if TYPE_CHECKING: @@ -97,11 +101,13 @@ def __init__( config: TelegramConfig, runtime: AdkRuntime, google_health_service: GoogleHealthService | None = None, + access_storage: TelegramAccessStorage | None = None, ) -> None: """Initialize the Telegram bot.""" self.config = config self.runtime = runtime self.google_health_service = google_health_service + self.access_storage = access_storage self._api: TelegramApiClient | None = None self._running = False self._polling_task: asyncio.Task[None] | None = None @@ -270,6 +276,9 @@ async def _polling_loop(self) -> None: async def _safe_handle_update(self, update: Update) -> None: """Handle update concurrently and allow cancellation.""" + if not await self._authorize_update(update): + return + if update.callback_query: # Handle callback queries immediately without cancelling conversation tasks try: @@ -318,6 +327,137 @@ async def _safe_handle_update(self, update: Update) -> None: conversation_key, current_seq, self._handle_update(update) ) + def _access_code_fingerprint(self) -> str: + """Return a non-plaintext marker used to invalidate rotated access codes.""" + access_code = self.config.telegram_access_code + if access_code is None: # pragma: no cover - guarded by caller + raise RuntimeError("Telegram access code is not configured") + return hashlib.sha256(access_code.encode("utf-8")).hexdigest() + + def _get_access_storage(self) -> TelegramAccessStorage: + """Get the injected or process-wide access storage.""" + return self.access_storage or get_telegram_access_storage() + + async def _authorize_update(self, update: Update) -> bool: + """Allow only authenticated private Telegram traffic into the bot.""" + if not self.config.access_control_enabled: + return True + + message = update.message + callback = update.callback_query + if message is None and callback is not None: + message = callback.message + if message is None or message.chat.type != ChatType.PRIVATE: + if callback is not None: + await self.api.answer_callback_query( + callback.id, text="Access required" + ) + return False + + sender = callback.from_user if callback is not None else message.from_user + if sender is None or sender.id != message.chat.id: + return False + + storage = self._get_access_storage() + fingerprint = self._access_code_fingerprint() + try: + authorized = await storage.is_authorized(sender.id, fingerprint) + has_authorization_record = await storage.has_authorization_record(sender.id) + if not authorized and not has_authorization_record: + authorized = await self._grant_legacy_access_if_applicable( + message, sender, storage + ) + if authorized: + await storage.record_identity(self._identity_from_user(sender)) + return True + if callback is not None: + await self.api.answer_callback_query( + callback.id, text="Access required" + ) + return False + return await self._handle_new_user_start(message, sender, storage) + except Exception: + logger.exception("Telegram access control failed closed") + return False + + async def _grant_legacy_access_if_applicable( + self, + message: Message, + sender: User, + storage: TelegramAccessStorage, + ) -> bool: + """Grandfather a historical direct chat without changing its session key.""" + session_identity = self._build_session_identity( + chat_id=str(message.chat.id), + message_thread_id=None, + ) + has_history = await self.runtime.has_existing_session( + locator=SessionLocator( + user_id=session_identity.user_id, + session_id_prefix=session_identity.session_id_prefix, + ) + ) + if not has_history: + return False + await storage.grant(sender.id, source="legacy") + return True + + async def _handle_new_user_start( + self, + message: Message, + sender: User, + storage: TelegramAccessStorage, + ) -> bool: + """Authenticate a new private user through the locally consumed /start code.""" + command, separator, supplied_code = (message.text or "").partition(" ") + if command != "/start" or not separator: + await self.api.send_message( + chat_id=message.chat.id, + text="Access required. Send /start followed by your access code.", + ) + return False + + configured_code = self.config.telegram_access_code + if configured_code is None: # pragma: no cover - guarded by caller + return False + is_valid = hmac.compare_digest(supplied_code.strip(), configured_code) + await self._delete_access_code_message(message) + if not is_valid: + await self.api.send_message( + chat_id=message.chat.id, + text="That access code is not valid. Please try again.", + ) + return False + + await storage.grant( + sender.id, + source="passphrase", + access_code_fingerprint=self._access_code_fingerprint(), + ) + await storage.record_identity(self._identity_from_user(sender)) + await self._send_start_message(message.chat.id) + return False + + async def _delete_access_code_message(self, message: Message) -> None: + """Best-effort removal of an access code from the visible private chat.""" + try: + await self.api.delete_message(message.chat.id, message.message_id) + except TelegramApiError: + logger.warning("Could not delete Telegram access-code message") + + def _identity_from_user(self, user: User) -> TelegramIdentity: + """Build a bounded local-only display label from Telegram profile fields.""" + display_name = " ".join( + part.strip() + for part in (user.first_name, user.last_name or "") + if part.strip() + )[:128] + return TelegramIdentity( + user_id=user.id, + display_name=display_name or "Telegram user", + username=user.username.strip()[:64] if user.username else None, + ) + async def _run_sequenced_turn( self, conversation_key: str, diff --git a/src/blacki/utils/config.py b/src/blacki/utils/config.py index 5af4b83..ac9c0ea 100644 --- a/src/blacki/utils/config.py +++ b/src/blacki/utils/config.py @@ -167,6 +167,12 @@ class ServerEnv(BaseModel): description="Telegram bot token from @BotFather", ) + telegram_access_code: str | None = Field( + default=None, + alias="TELEGRAM_ACCESS_CODE", + description="Shared access code required for new private Telegram users", + ) + model_config = ConfigDict( populate_by_name=True, # Allow both field names and aliases extra="ignore", # Ignore extra env vars (system vars, etc.) @@ -190,6 +196,8 @@ def print_config(self) -> None: print(f"TELEGRAM_ENABLED: {self.telegram_enabled}") masked_token = "********" if self.telegram_bot_token else "[not set]" print(f"TELEGRAM_BOT_TOKEN: {masked_token}") + masked_access_code = "********" if self.telegram_access_code else "[not set]" + print(f"TELEGRAM_ACCESS_CODE: {masked_access_code}") print() @property diff --git a/tests/dashboard/test_data.py b/tests/dashboard/test_data.py index 52e3c4b..8d58e16 100644 --- a/tests/dashboard/test_data.py +++ b/tests/dashboard/test_data.py @@ -25,6 +25,7 @@ _iter_raw_jsonl_lines, _JsonlScan, _load_snapshot_sync, + _load_telegram_identities_sync, _log_item, _normalize_epoch_seconds, _number, @@ -118,6 +119,28 @@ def _insert_event( connection.close() +def _make_identity_db(path: Path) -> None: + connection = sqlite3.connect(path) + connection.execute( + """ + CREATE TABLE telegram_identities ( + telegram_user_id INTEGER PRIMARY KEY, + display_name TEXT NOT NULL, + username TEXT, + updated_at TEXT NOT NULL + ) + """ + ) + connection.execute( + """ + INSERT INTO telegram_identities VALUES (?, ?, ?, ?) + """, + (42, "Ada Lovelace", "ada", "2026-08-20T00:00:00+00:00"), + ) + connection.commit() + connection.close() + + def _text_event( event_id: str, role: str, @@ -235,6 +258,7 @@ async def test_sessions_users_replay_versions_and_overview(tmp_path: Path) -> No "human", "model", ] + assert replay["messages"][0]["text"] == "" assert [ message["mime_type"] @@ -258,6 +282,34 @@ async def test_sessions_users_replay_versions_and_overview(tmp_path: Path) -> No assert overview["activity"] +@pytest.mark.asyncio +async def test_dashboard_labels_private_telegram_users_from_local_identity_db( + tmp_path: Path, +) -> None: + session_db = tmp_path / "sessions.db" + identity_db = tmp_path / "tools.db" + _make_db(session_db) + _insert_session(session_db, "telegram-chat-42", "telegram-chat-42-v1") + _make_identity_db(identity_db) + + users = await DashboardStore( + session_db, tmp_path, identity_db_path=identity_db + ).list_users("ada", 50, 0) + + assert users["total"] == 1 + assert users["items"][0]["display_name"] == "Ada Lovelace" + assert users["items"][0]["username"] == "ada" + + +def test_telegram_identity_reader_tolerates_missing_and_invalid_databases( + tmp_path: Path, +) -> None: + assert _load_telegram_identities_sync(tmp_path / "missing.db") == {} + corrupt = tmp_path / "corrupt.db" + corrupt.write_text("not a sqlite database", encoding="utf-8") + assert _load_telegram_identities_sync(corrupt) == {} + + @pytest.mark.asyncio async def test_overview_error_rate_uses_failed_invocations_only(tmp_path: Path) -> None: db_path = tmp_path / "sessions.db" diff --git a/tests/dashboard/test_routes.py b/tests/dashboard/test_routes.py index ef0e096..403668b 100644 --- a/tests/dashboard/test_routes.py +++ b/tests/dashboard/test_routes.py @@ -235,7 +235,12 @@ def test_default_store_uses_session_db_log_dir_and_app_name(tmp_path: Path) -> N client, _ = _client(tmp_path, use_default_store=True) assert client.get("/dashboard/api/overview").status_code == 200 - ctor.assert_called_once_with(tmp_path / ".adk" / "sessions.db", log_dir, "blacki") + ctor.assert_called_once_with( + tmp_path / ".adk" / "sessions.db", + log_dir, + "blacki", + tmp_path / ".adk" / "tools.db", + ) def test_real_store_degrades_cleanly_when_local_records_are_missing( diff --git a/tests/test_adk_runtime.py b/tests/test_adk_runtime.py index c6e3c87..302492c 100644 --- a/tests/test_adk_runtime.py +++ b/tests/test_adk_runtime.py @@ -193,6 +193,19 @@ async def test_create_next_session_increments_version() -> None: assert second_session.id == "telegram-chat-123-v2" +@pytest.mark.asyncio +async def test_has_existing_session_does_not_create_a_new_session() -> None: + runtime = AdkRuntime(InMemorySessionService()) + locator = SessionLocator( + user_id="telegram-chat-123", + session_id_prefix="telegram-chat-123", + ) + + assert await runtime.has_existing_session(locator=locator) is False + await runtime.create_next_session(locator=locator) + assert await runtime.has_existing_session(locator=locator) is True + + async def test_get_or_create_session_reuses_latest_version() -> None: """Test that the active session resolves to the latest version.""" runtime = AdkRuntime(InMemorySessionService()) diff --git a/tests/test_container.py b/tests/test_container.py index a5969a1..42c7afd 100644 --- a/tests/test_container.py +++ b/tests/test_container.py @@ -195,6 +195,16 @@ async def test_preferences_storage_property(self, conn, lock) -> None: assert storage is not None assert container._preferences_storage is storage + @pytest.mark.asyncio + async def test_telegram_access_storage_property(self, conn, lock) -> None: + """Should lazily instantiate Telegram access storage.""" + container = AppContainer(conn=conn, _lock=lock) + + storage = container.telegram_access_storage + + assert storage is not None + assert container._telegram_access_storage is storage + @pytest.mark.asyncio async def test_declarative_db_storage_property(self, conn, lock) -> None: """Should lazily instantiate declarative DB storage.""" diff --git a/tests/test_server_config.py b/tests/test_server_config.py index f2b353d..9acfd72 100644 --- a/tests/test_server_config.py +++ b/tests/test_server_config.py @@ -28,6 +28,7 @@ def mock_dependencies() -> Generator[MagicMock]: mock_env.serve_web_interface = True mock_env.reload_agents = False mock_env.sqlite_path = None + mock_env.telegram_access_code = None mock_env.agent_dir = "src" mock_env.host = "127.0.0.1" diff --git a/tests/test_telegram_access.py b/tests/test_telegram_access.py new file mode 100644 index 0000000..ed49601 --- /dev/null +++ b/tests/test_telegram_access.py @@ -0,0 +1,293 @@ +"""Tests for local Telegram authorization and identity storage.""" + +import asyncio +from collections.abc import AsyncIterator +from pathlib import Path +from typing import Any, cast +from unittest.mock import AsyncMock, create_autospec, patch + +import pytest + +from blacki.storage.sqlite import create_connection +from blacki.telegram import TelegramConfig +from blacki.telegram.access import ( + TelegramAccessStorage, + TelegramIdentity, + get_telegram_access_storage, +) +from blacki.telegram.api import TelegramApiClient, TelegramApiError +from blacki.telegram.bot import TelegramBot +from blacki.telegram.types import Update + + +@pytest.fixture +async def storage(tmp_path: Path) -> AsyncIterator[TelegramAccessStorage]: + connection = await create_connection(tmp_path / "tools.db") + result = TelegramAccessStorage(connection, asyncio.Lock()) + await result.initialize() + try: + yield result + finally: + await connection.close() + + +@pytest.mark.asyncio +async def test_passphrase_authorization_is_invalidated_by_rotation( + storage: TelegramAccessStorage, +) -> None: + await storage.grant( + 42, + source="passphrase", + access_code_fingerprint="before-rotation", + ) + + assert await storage.is_authorized(42, "before-rotation") is True + assert await storage.is_authorized(42, "after-rotation") is False + assert await storage.has_authorization_record(42) is True + + +@pytest.mark.asyncio +async def test_legacy_authorization_survives_access_code_rotation( + storage: TelegramAccessStorage, +) -> None: + await storage.grant(42, source="legacy") + + assert await storage.is_authorized(42, "before-rotation") is True + assert await storage.is_authorized(42, "after-rotation") is True + + +@pytest.mark.asyncio +async def test_identity_is_updated_without_affecting_authorization( + storage: TelegramAccessStorage, +) -> None: + await storage.grant(42, source="legacy") + await storage.record_identity(TelegramIdentity(42, "First Name", "first")) + await storage.record_identity(TelegramIdentity(42, "New Name", "new")) + + row = await storage._fetch_one( + """ + SELECT display_name, username + FROM telegram_identities + WHERE telegram_user_id = ? + """, + (42,), + ) + + assert row == {"display_name": "New Name", "username": "new"} + assert await storage.is_authorized(42, "any-code") is True + + +@pytest.mark.asyncio +async def test_get_telegram_access_storage_uses_initialized_container( + tmp_path: Path, +) -> None: + from blacki.container import ( + reset_container_for_tests, + set_container_from_connection, + ) + + connection = await create_connection(tmp_path / "tools.db") + container = set_container_from_connection(connection) + try: + with pytest.raises(RuntimeError, match="not initialized"): + get_telegram_access_storage() + await container.telegram_access_storage.initialize() + assert get_telegram_access_storage() is container.telegram_access_storage + finally: + reset_container_for_tests() + await connection.close() + + +class _Runtime: + def __init__(self, has_history: bool) -> None: + self.has_history = has_history + + async def has_existing_session(self, **_kwargs: object) -> bool: + return self.has_history + + +def _update(text: str, *, chat_type: str = "private", sender_id: int = 42) -> Update: + return Update.model_validate( + { + "update_id": 1, + "message": { + "message_id": 2, + "date": "2026-08-20T00:00:00Z", + "chat": {"id": 42, "type": chat_type}, + "from": { + "id": sender_id, + "first_name": "Ada", + "username": "ada", + }, + "text": text, + }, + } + ) + + +def _callback_update(*, message: dict[str, object] | None) -> Update: + return Update.model_validate( + { + "update_id": 1, + "callback_query": { + "id": "callback-1", + "from": {"id": 42, "first_name": "Ada"}, + "chat_instance": "instance-1", + "data": "setting:model", + "message": message, + }, + } + ) + + +def _bot(storage: TelegramAccessStorage, *, has_history: bool) -> TelegramBot: + config = TelegramConfig.model_validate( + { + "TELEGRAM_ENABLED": True, + "TELEGRAM_BOT_TOKEN": "test-token-123", + "TELEGRAM_ACCESS_CODE": "test-access-code", + } + ) + bot = TelegramBot(config, _Runtime(has_history), access_storage=storage) # type: ignore[arg-type] + api = create_autospec(TelegramApiClient, instance=True) + api.send_message = AsyncMock() + api.delete_message = AsyncMock() + bot._api = api + return bot + + +@pytest.mark.asyncio +async def test_start_access_code_authorizes_new_user( + storage: TelegramAccessStorage, +) -> None: + bot = _bot(storage, has_history=False) + + allowed = await bot._authorize_update(_update("/start test-access-code")) + + assert allowed is False + assert await storage.is_authorized(42, bot._access_code_fingerprint()) is True + assert cast(AsyncMock, cast(Any, bot.api).delete_message).await_count == 1 + + +@pytest.mark.asyncio +async def test_historical_private_chat_is_grandfathered( + storage: TelegramAccessStorage, +) -> None: + bot = _bot(storage, has_history=True) + + allowed = await bot._authorize_update(_update("normal message")) + + assert allowed is True + assert await storage.is_authorized(42, "rotated-code") is True + + +@pytest.mark.asyncio +async def test_rotated_passphrase_user_is_not_regrandfathered( + storage: TelegramAccessStorage, +) -> None: + bot = _bot(storage, has_history=True) + await storage.grant( + 42, + source="passphrase", + access_code_fingerprint="old-code", + ) + + assert await bot._authorize_update(_update("normal message")) is False + assert await storage.is_authorized(42, bot._access_code_fingerprint()) is False + + +@pytest.mark.asyncio +async def test_authorized_user_bypasses_legacy_lookup( + storage: TelegramAccessStorage, +) -> None: + bot = _bot(storage, has_history=False) + await storage.grant( + 42, + source="passphrase", + access_code_fingerprint=bot._access_code_fingerprint(), + ) + + assert await bot._authorize_update(_update("normal message")) is True + + +@pytest.mark.asyncio +async def test_new_user_without_start_code_is_denied( + storage: TelegramAccessStorage, +) -> None: + bot = _bot(storage, has_history=False) + + assert await bot._authorize_update(_update("hello")) is False + cast(Any, bot.api).send_message.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_invalid_start_code_is_deleted_and_denied( + storage: TelegramAccessStorage, +) -> None: + bot = _bot(storage, has_history=False) + api = cast(Any, bot.api) + api.delete_message.side_effect = TelegramApiError("delete failed", 400) + + assert await bot._authorize_update(_update("/start wrong-code")) is False + assert await storage.is_authorized(42, bot._access_code_fingerprint()) is False + api.delete_message.assert_awaited_once() + api.send_message.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_group_and_mismatched_sender_are_rejected( + storage: TelegramAccessStorage, +) -> None: + bot = _bot(storage, has_history=False) + + assert await bot._authorize_update(_update("hello", chat_type="group")) is False + assert await bot._authorize_update(_update("hello", sender_id=7)) is False + + +@pytest.mark.asyncio +async def test_callback_without_private_message_is_rejected( + storage: TelegramAccessStorage, +) -> None: + bot = _bot(storage, has_history=False) + + assert await bot._authorize_update(_callback_update(message=None)) is False + cast(Any, bot.api).answer_callback_query.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_unauthorized_callback_is_rejected_after_identity_check( + storage: TelegramAccessStorage, +) -> None: + bot = _bot(storage, has_history=False) + message = { + "message_id": 2, + "date": "2026-08-20T00:00:00Z", + "chat": {"id": 42, "type": "private"}, + "from": {"id": 42, "first_name": "Ada"}, + "text": "Settings", + } + + assert await bot._authorize_update(_callback_update(message=message)) is False + cast(Any, bot.api).answer_callback_query.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_access_storage_failure_fails_closed( + storage: TelegramAccessStorage, +) -> None: + bot = _bot(storage, has_history=False) + cast(Any, storage).is_authorized = AsyncMock(side_effect=RuntimeError("database")) + + assert await bot._authorize_update(_update("hello")) is False + + +@pytest.mark.asyncio +async def test_safe_handler_does_not_route_denied_update( + storage: TelegramAccessStorage, +) -> None: + bot = _bot(storage, has_history=False) + handle_update = AsyncMock() + with patch.object(bot, "_handle_update", new=handle_update): + await bot._safe_handle_update(_update("hello")) + + handle_update.assert_not_awaited()