add password authentication for bot users

README.md

@@ -39,9 +39,14 @@ python main.py
 
 ```dotenv
 BOT_TOKEN=...
+ASSISTANT_PASSWORD=...
 ASSISTANT_MODE=local
 ```
 
+`ASSISTANT_PASSWORD` — общий пароль доступа к боту. Новый пользователь должен
+один раз отправить его боту в личном чате; после успешной проверки авторизация
+сохраняется в SQLite, а сам пароль в базу данных не записывается.
+
 `ASSISTANT_MODE` принимает `local` (значение по умолчанию) или `yandex`.
 Общие дополнительные настройки: `ASSISTANT_DB`, `ASSISTANT_TIMEZONE` и
 `VOICE_MAX_DURATION_SECONDS`.

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,

tests/test_auth.py

@@ -0,0 +1,75 @@
+import tempfile
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
+from unittest.mock import AsyncMock
+
+from telegram.constants import ChatType
+from telegram.ext import ApplicationHandlerStop
+
+from assistant_bot.auth import authentication_guard
+from assistant_bot.storage import AssistantStorage
+
+
+class AuthenticationGuardTests(unittest.IsolatedAsyncioTestCase):
+    def setUp(self) -> None:
+        self.temporary_directory = tempfile.TemporaryDirectory()
+        self.storage = AssistantStorage(
+            Path(self.temporary_directory.name) / "assistant.sqlite3"
+        )
+
+    def tearDown(self) -> None:
+        self.temporary_directory.cleanup()
+
+    def make_update_and_context(self, text: str) -> tuple[Any, Any]:
+        message = SimpleNamespace(text=text, reply_text=AsyncMock())
+        update = SimpleNamespace(
+            effective_user=SimpleNamespace(id=42),
+            effective_chat=SimpleNamespace(type=ChatType.PRIVATE),
+            effective_message=message,
+            inline_query=None,
+        )
+        context = SimpleNamespace(
+            application=SimpleNamespace(
+                bot_data={
+                    "storage": self.storage,
+                    "assistant_password": "секрет",
+                }
+            )
+        )
+        return update, context
+
+    async def test_correct_password_authorizes_user(self) -> None:
+        update, context = self.make_update_and_context("секрет")
+
+        with self.assertRaises(ApplicationHandlerStop):
+            await authentication_guard(update, context)
+
+        self.assertTrue(self.storage.is_user_authorized(42))
+        update.effective_message.reply_text.assert_awaited_once_with(
+            "Пароль принят. Доступ открыт — повторно вводить его не нужно."
+        )
+
+    async def test_wrong_password_does_not_authorize_user(self) -> None:
+        update, context = self.make_update_and_context("неверно")
+
+        with self.assertRaises(ApplicationHandlerStop):
+            await authentication_guard(update, context)
+
+        self.assertFalse(self.storage.is_user_authorized(42))
+        update.effective_message.reply_text.assert_awaited_once_with(
+            "Неверный пароль. Попробуй еще раз."
+        )
+
+    async def test_authorized_user_passes_guard(self) -> None:
+        self.storage.authorize_user(42)
+        update, context = self.make_update_and_context("обычное сообщение")
+
+        await authentication_guard(update, context)
+
+        update.effective_message.reply_text.assert_not_awaited()
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_config.py

@@ -67,5 +67,25 @@ class AssistantModeConfigTests(unittest.TestCase):
             )
 
 
+class PasswordConfigTests(unittest.TestCase):
+    def test_password_is_required(self) -> None:
+        with patch("assistant_bot.config.load_env_file"), patch.dict(
+            os.environ,
+            {},
+            clear=False,
+        ):
+            os.environ.pop(config.PASSWORD_ENV_NAME, None)
+            with self.assertRaises(RuntimeError):
+                config.get_assistant_password()
+
+    def test_password_is_read_from_environment(self) -> None:
+        with patch("assistant_bot.config.load_env_file"), patch.dict(
+            os.environ,
+            {config.PASSWORD_ENV_NAME: "  test-password  "},
+            clear=False,
+        ):
+            self.assertEqual(config.get_assistant_password(), "test-password")
+
+
 if __name__ == "__main__":
     unittest.main()

tests/test_core.py

@@ -53,6 +53,18 @@ class AgentDecisionTests(unittest.TestCase):
 
 
 class StorageTests(unittest.TestCase):
+    def test_user_authorization_is_persisted(self) -> None:
+        with tempfile.TemporaryDirectory() as directory:
+            database_path = Path(directory) / "assistant.sqlite3"
+            storage = AssistantStorage(database_path)
+
+            self.assertFalse(storage.is_user_authorized(42))
+            storage.authorize_user(42)
+
+            reopened_storage = AssistantStorage(database_path)
+            self.assertTrue(reopened_storage.is_user_authorized(42))
+            self.assertFalse(reopened_storage.is_user_authorized(43))
+
     def test_legacy_user_settings_gets_yandex_model_column(self) -> None:
         with tempfile.TemporaryDirectory() as directory:
             database_path = Path(directory) / "assistant.sqlite3"