add password authentication for bot users

assistant_bot/application.py

@@ -6,12 +6,15 @@ from telegram.ext import (
     CommandHandler,
     InlineQueryHandler,
     MessageHandler,
+    TypeHandler,
     filters,
 )
 
 from .ai import AIClient
+from .auth import authentication_guard
 from .config import (
     get_assistant_mode,
+    get_assistant_password,
     get_bot_token,
     get_db_path,
     get_local_timezone,
@@ -53,6 +56,7 @@ def configure_logging() -> None:
 
 
 def create_application() -> Application:
+    assistant_password = get_assistant_password()
     storage = AssistantStorage(get_db_path())
     mode = get_assistant_mode()
     application = (
@@ -63,6 +67,7 @@ def create_application() -> Application:
         .build()
     )
     application.bot_data["storage"] = storage
+    application.bot_data["assistant_password"] = assistant_password
     ai_client: AIClient
     speech_recognizer: SpeechTranscriber
     if mode == "local":
@@ -88,6 +93,7 @@ def create_application() -> Application:
     application.bot_data["speech_recognizer"] = speech_recognizer
     application.bot_data["voice_max_duration"] = get_voice_max_duration_seconds()
 
+    application.add_handler(TypeHandler(Update, authentication_guard), group=-1)
     application.add_handler(CommandHandler("start", start))
     application.add_handler(CommandHandler("help", help_command))
     application.add_handler(CommandHandler("ask", ask_command))

assistant_bot/auth.py

@@ -0,0 +1,57 @@
+from hmac import compare_digest
+
+from telegram import Update
+from telegram.constants import ChatType
+from telegram.ext import ApplicationHandlerStop, ContextTypes
+
+from .telegram_utils import get_storage, require_user_id
+
+
+def passwords_match(candidate: str, expected: str) -> bool:
+    return compare_digest(
+        candidate.encode("utf-8"),
+        expected.encode("utf-8"),
+    )
+
+
+async def authentication_guard(
+    update: Update,
+    context: ContextTypes.DEFAULT_TYPE,
+) -> None:
+    user_id = require_user_id(update)
+    if user_id is None:
+        raise ApplicationHandlerStop
+
+    storage = get_storage(context)
+    if storage.is_user_authorized(user_id):
+        return
+
+    message = update.effective_message
+    chat = update.effective_chat
+    if chat is not None and chat.type != ChatType.PRIVATE:
+        if message:
+            await message.reply_text(
+                "Сначала авторизуйся в личном чате с ботом."
+            )
+        raise ApplicationHandlerStop
+
+    candidate = message.text.strip() if message and message.text else ""
+    password = str(context.application.bot_data["assistant_password"])
+    if candidate and passwords_match(candidate, password):
+        storage.authorize_user(user_id)
+        await message.reply_text(
+            "Пароль принят. Доступ открыт — повторно вводить его не нужно."
+        )
+        raise ApplicationHandlerStop
+
+    if candidate and not candidate.startswith("/"):
+        await message.reply_text("Неверный пароль. Попробуй еще раз.")
+        raise ApplicationHandlerStop
+
+    if message:
+        await message.reply_text(
+            "Для доступа к боту введи пароль одним текстовым сообщением."
+        )
+    elif update.inline_query:
+        await update.inline_query.answer([], cache_time=0)
+    raise ApplicationHandlerStop

assistant_bot/config.py

@@ -10,6 +10,7 @@ PROJECT_ROOT = Path(__file__).resolve().parent.parent
 ENV_FILE = PROJECT_ROOT / ".env"
 
 TOKEN_ENV_NAME = "BOT_TOKEN"
+PASSWORD_ENV_NAME = "ASSISTANT_PASSWORD"
 DB_ENV_NAME = "ASSISTANT_DB"
 MODE_ENV_NAME = "ASSISTANT_MODE"
 OLLAMA_BASE_URL_ENV_NAME = "OLLAMA_BASE_URL"
@@ -80,6 +81,16 @@ def get_bot_token() -> str:
     return token
 
 
+def get_assistant_password() -> str:
+    load_env_file()
+    password = os.getenv(PASSWORD_ENV_NAME, "").strip()
+    if not password:
+        raise RuntimeError(
+            f"Set {PASSWORD_ENV_NAME} in environment or .env file."
+        )
+    return password
+
+
 def get_db_path() -> Path:
     load_env_file()
     raw_path = os.getenv(DB_ENV_NAME)

assistant_bot/storage.py

@@ -55,6 +55,11 @@ class AssistantStorage:
                     updated_at TEXT NOT NULL
                 );
 
+                CREATE TABLE IF NOT EXISTS authorized_users (
+                    user_id INTEGER PRIMARY KEY,
+                    authorized_at TEXT NOT NULL
+                );
+
                 CREATE TABLE IF NOT EXISTS memories (
                     id INTEGER PRIMARY KEY AUTOINCREMENT,
                     user_id INTEGER NOT NULL,
@@ -127,6 +132,25 @@ class AssistantStorage:
                     "ALTER TABLE user_settings ADD COLUMN yandex_model TEXT"
                 )
 
+    def is_user_authorized(self, user_id: int) -> bool:
+        with self._connection() as connection:
+            row = connection.execute(
+                "SELECT 1 FROM authorized_users WHERE user_id = ?",
+                (user_id,),
+            ).fetchone()
+        return row is not None
+
+    def authorize_user(self, user_id: int) -> None:
+        with self._connection() as connection:
+            connection.execute(
+                """
+                INSERT INTO authorized_users(user_id, authorized_at)
+                VALUES (?, ?)
+                ON CONFLICT(user_id) DO NOTHING
+                """,
+                (user_id, to_utc_iso(utc_now())),
+            )
+
     def add_conversation_exchange(
         self,
         user_id: int,