From 053b124a9a2ac001da3bf98514cb2c90139d327a Mon Sep 17 00:00:00 2001 From: kandrusyak Date: Sat, 25 Jul 2026 15:57:30 +0300 Subject: [PATCH] add password authentication for bot users --- README.md | 5 +++ assistant_bot/application.py | 6 +++ assistant_bot/auth.py | 57 +++++++++++++++++++++++++++ assistant_bot/config.py | 11 ++++++ assistant_bot/storage.py | 24 ++++++++++++ tests/test_auth.py | 75 ++++++++++++++++++++++++++++++++++++ tests/test_config.py | 20 ++++++++++ tests/test_core.py | 12 ++++++ 8 files changed, 210 insertions(+) create mode 100644 assistant_bot/auth.py create mode 100644 tests/test_auth.py diff --git a/README.md b/README.md index 52d8919..ebf0634 100644 --- a/README.md +++ b/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`. diff --git a/assistant_bot/application.py b/assistant_bot/application.py index fdf705d..79ff43d 100644 --- a/assistant_bot/application.py +++ b/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)) diff --git a/assistant_bot/auth.py b/assistant_bot/auth.py new file mode 100644 index 0000000..2807285 --- /dev/null +++ b/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 diff --git a/assistant_bot/config.py b/assistant_bot/config.py index b383ce3..76086e7 100644 --- a/assistant_bot/config.py +++ b/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) diff --git a/assistant_bot/storage.py b/assistant_bot/storage.py index 43c5670..b21f2b3 100644 --- a/assistant_bot/storage.py +++ b/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, diff --git a/tests/test_auth.py b/tests/test_auth.py new file mode 100644 index 0000000..eb568dc --- /dev/null +++ b/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() diff --git a/tests/test_config.py b/tests/test_config.py index 7b37f00..9005ccf 100644 --- a/tests/test_config.py +++ b/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() diff --git a/tests/test_core.py b/tests/test_core.py index 05c6436..49b4aa3 100644 --- a/tests/test_core.py +++ b/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"