import asyncio from pathlib import Path import sys import tempfile import unittest from unittest.mock import AsyncMock, patch ROOT = Path(__file__).resolve().parents[1] sys.path[:0] = [str(ROOT), str(ROOT / "vendor")] from awg_gus.bot import Bot, Event from awg_gus.core import Awg, AwgError, State, validate_command from awg_gus.main import pair, validate_url OWNER = "11111111-1111-4111-8111-111111111111" OTHER = "22222222-2222-4222-8222-222222222222" CHAT = "33333333-3333-4333-8333-333333333333" MSG = "44444444-4444-4444-8444-444444444444" class FakeAwg: def __init__(self): self.calls = [] self.rows = [{"name": "alice", "ip": "10.8.1.2", "online": True, "handshake": 1, "ago": 10, "rx": 1024, "tx": 512}] async def call(self, *args, **kw): validate_command(args) self.calls.append(args) data = {} if args[:2] == ("clients", "list"): data = self.rows elif args[:2] == ("client", "conf"): data = {"text": "[Interface]\nPrivateKey = test-key\n", "file": "/tmp/alice.conf"} elif args[:2] == ("client", "add"): self.rows.append({"name": args[2], "ip": "10.8.1.3"}) elif args[:2] == ("client", "del"): self.rows[:] = [r for r in self.rows if r["name"] != args[2]] elif args == ("status",): data = {"version": "test", "host": "test-host", "server": {"up": True, "clients": len(self.rows)}} elif args == ("server", "info"): data = {"up": True} return {"ok": True, "data": data, "log": "done"} class FakeDelivery: def __init__(self): self.sent = [] self.acks = [] async def show(self, chat, text, buttons=None, message=None): self.sent.append(dict(chat=chat, text=text, buttons=buttons, message=message)) return {"message_id": MSG} async def ack(self, cid, text=None): self.acks.append((cid, text)) async def call(self, method, **params): self.sent.append({"method": method, **params}) return {"message_id": MSG} def message(n, text, uid=OWNER, kind="dm"): return {"update_id": n, "message": {"message_id": MSG, "text": text, "from": {"id": uid}, "chat": {"id": CHAT, "type": kind}}} def callback(n, data, uid=OWNER, kind="dm"): m = message(n, "", uid, kind)["message"] return {"update_id": n, "callback_query": {"id": f"cb-{n}", "from": {"id": uid}, "data": data, "message": m}} class BotTests(unittest.IsolatedAsyncioTestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() self.path = Path(self.tmp.name) / "state.sqlite3" self.state = State(self.path, "test") self.api, self.io = FakeAwg(), FakeDelivery() self.bot = Bot({"owners": [OWNER]}, self.api, self.io, self.state) async def asyncTearDown(self): for t in self.bot.tasks.values(): t.cancel() await asyncio.gather(*self.bot.tasks.values(), return_exceptions=True) self.state.db.close() self.tmp.cleanup() async def feed(self, u): await self.bot.process(u) await asyncio.gather(*self.bot.tasks.values()) async def test_no_access_for_other_users_or_channels(self): await self.feed(message(1, "/menu", OTHER)) await self.feed(message(2, "/menu", OWNER, "channel")) await self.feed(callback(3, "delete:alice", OTHER)) await self.feed(callback(4, "conf:alice", OWNER, "channel")) self.assertEqual(self.api.calls, []) self.assertEqual(self.io.sent, []) self.assertEqual(len(self.io.acks), 2) async def test_owner_menu_and_callback_ack(self): await self.feed(message(1, "/start")) await self.feed(callback(2, "list:0")) self.assertIn(("status",), self.api.calls) self.assertIn(("clients", "list"), self.api.calls) self.assertEqual(self.io.acks[0], ("cb-2", None)) self.assertIn("alice", str(self.io.sent)) async def test_delete_requires_confirmation_and_is_single_use(self): await self.feed(callback(1, "delete:alice")) self.assertNotIn(("client", "del", "alice"), self.api.calls) token = next(iter(self.bot.confirmations)) await self.feed(callback(2, "yes:" + token, OTHER)) self.assertIn(token, self.bot.confirmations) await self.feed(callback(3, "yes:" + token)) await self.feed(callback(4, "yes:" + token)) self.assertEqual(self.api.calls.count(("client", "del", "alice")), 1) async def test_cancel_invalidates_confirmation(self): await self.feed(callback(1, "delete:alice")) token = next(iter(self.bot.confirmations)) await self.feed(message(2, "/cancel")) await self.feed(callback(3, "yes:" + token)) self.assertNotIn(("client", "del", "alice"), self.api.calls) async def test_expired_confirmation_does_not_execute(self): await self.feed(callback(1, "delete:alice")) token = next(iter(self.bot.confirmations)) v = self.bot.confirmations[token] self.bot.confirmations[token] = (*v[:2], 0, *v[3:]) await self.feed(callback(2, "yes:" + token)) self.assertNotIn(("client", "del", "alice"), self.api.calls) async def test_create_flow_delivers_config_without_leaking_in_text(self): with patch("awg_gus.bot.shutil.which", return_value=None): await self.feed(message(1, "/new")) await self.feed(message(2, "bob")) await self.feed(callback(3, "newexp:+30d")) token = next(iter(self.bot.confirmations)) await self.feed(callback(4, "yes:" + token)) self.assertIn(("client", "add", "bob", "mimicry=server", "expire=+30d"), self.api.calls) sent = [m for m in self.io.sent if m.get("files")] self.assertEqual(len(sent), 1) self.assertEqual(sent[0]["files"][0].read()[1], "bob.conf") self.assertNotIn("test-key", " ".join(s.get("text", "") for s in self.io.sent)) async def test_invalid_name_never_reaches_awg(self): await self.feed(message(1, "/new")) await self.feed(message(2, "alice; reboot")) await self.feed(callback(3, "conf:../../passwd")) self.assertFalse(any(a[:2] in {("client", "add"), ("client", "conf")} for a in self.api.calls)) async def test_limit_speed_and_expiry_use_validated_commands(self): n = 0 for action, value, expected in [ ("limit", "50G total", ("client", "limit", "alice", "50G", "total")), ("speed", "20", ("client", "speed", "alice", "20"))]: n += 1 await self.feed(callback(n, action + ":alice")) n += 1 await self.feed(message(n, value)) n += 1 await self.feed(callback(n, "yes:" + next(iter(self.bot.confirmations)))) self.assertIn(expected, self.api.calls) await self.feed(callback(n+1, "exp:alice:+7d")) await self.feed(callback(n+2, "yes:" + next(iter(self.bot.confirmations)))) self.assertIn(("client", "expire", "alice", "+7d"), self.api.calls) async def test_duplicate_update_stays_deduplicated_after_restart(self): await self.feed(callback(1, "list:0")) count = len(self.api.calls) self.state.db.close() self.state = State(self.path, "test") self.bot = Bot({"owners": [OWNER]}, self.api, self.io, self.state) await self.feed(callback(1, "list:0")) self.assertEqual(len(self.api.calls), count) async def test_busy_callback_is_acknowledged_without_second_operation(self): gate = asyncio.Event() task = asyncio.create_task(gate.wait()) self.bot.tasks[(OWNER, CHAT)] = task await self.bot.process(callback(1, "delete:alice")) self.assertIn("выполняется", self.io.acks[-1][1]) self.assertEqual(self.api.calls, []) gate.set() await task async def test_interrupted_operation_is_preserved_on_cancellation(self): gate = asyncio.Event() async def blocked(*args, **kwargs): await gate.wait() self.bot.home = blocked await self.bot.process(message(1, "/start")) await asyncio.sleep(0) task = self.bot.tasks[(OWNER, CHAT)] task.cancel() await asyncio.gather(task, return_exceptions=True) self.assertEqual(self.state.db.execute("SELECT COUNT(*) FROM pending").fetchone()[0], 1) async def test_monitor_silent_until_state_changes_and_uses_chat_id(self): self.state.claim(1, OWNER, CHAT) self.state.finish(1) self.state.watch(OWNER, "alice") await self.bot.monitor_tick() self.assertEqual(self.io.sent, []) self.api.rows[0]["ago"] = 500 await self.bot.monitor_tick() await self.bot.monitor_tick() self.assertEqual(len(self.io.sent), 1) self.assertEqual(self.io.sent[0]["chat"], CHAT) self.assertIn("офлайн", self.io.sent[0]["text"]) async def test_config_and_qr_sent_in_same_message(self): proc = AsyncMock() proc.returncode = 0 proc.communicate.return_value = (b"\x89PNG\r\n\x1a\nmock", None) with patch("awg_gus.bot.shutil.which", return_value="qrencode"), patch( "awg_gus.bot.asyncio.create_subprocess_exec", return_value=proc) as spawn: await self.bot.send_config(Event(OWNER, CHAT, 1), "alice") self.assertEqual(len(self.io.sent[-1]["files"]), 2) self.assertNotIn("test-key", str(spawn.call_args)) self.assertIn(b"test-key", proc.communicate.call_args.args[0]) async def test_cancelled_qr_is_reaped_without_sending_files(self): from unittest.mock import Mock proc = AsyncMock() proc.kill = Mock() proc.communicate.side_effect = asyncio.CancelledError with patch("awg_gus.bot.shutil.which", return_value="qrencode"), patch( "awg_gus.bot.asyncio.create_subprocess_exec", return_value=proc): with self.assertRaises(asyncio.CancelledError): await self.bot.send_config(Event(OWNER, CHAT, 1), "alice") proc.kill.assert_called_once() proc.wait.assert_awaited_once() self.assertEqual(self.io.sent, []) class CoreTests(unittest.TestCase): def test_command_allowlist_rejects_root_escape_routes(self): for command in [("server", "reset"), ("bot", "update"), ("job", "start", "server", "reboot"), ("client", "add", "ok", "mimicry=server", "expire=$(id)"), ("client", "conf", "../root"), ("client", "speed", "alice", "1;id"), ("backup", "restore", "/tmp/x")]: with self.subTest(command=command), self.assertRaises(ValueError): validate_command(command) def test_url_rejects_insecure_or_embedded_credentials(self): for u in ["http://example.com", "https://user:pass@example.com", "https://example.com?x=1", "https://example.com/api/bot/v1"]: with self.assertRaises(ValueError): validate_url(u) self.assertEqual(validate_url("https://example.com/api/"), "https://example.com") def test_pairing_does_not_trust_first_update_or_group_code(self): from unittest.mock import Mock client = Mock() client.get_updates.side_effect = [[message(1, "Привет, бот!", OTHER), message(2, "/start SECRET", OTHER, "channel"), message(3, "/start wrong", OTHER), message(4, "/start SECRET")], []] self.assertEqual(pair(client, "SECRET"), (OWNER, CHAT, 5)) def test_state_cannot_be_shared_between_different_bots(self): with tempfile.TemporaryDirectory() as d: s = State(Path(d) / "s.db", "one") s.db.close() with self.assertRaises(ValueError): State(Path(d) / "s.db", "two") class ProcessTests(unittest.IsolatedAsyncioTestCase): async def test_json_api_without_shell_or_secret_arguments(self): proc = AsyncMock() proc.returncode = 0 proc.communicate.return_value = (b'progress\n{"ok":true,"data":{"version":"test"}}\n', b'') with patch("awg_gus.core.asyncio.create_subprocess_exec", return_value=proc) as spawn: result = await Awg("/usr/local/bin/awg2").call("version") self.assertEqual(result["data"]["version"], "test") self.assertEqual(spawn.call_args.args, ("/usr/local/bin/awg2", "api", "version")) async def test_malformed_or_failed_api_not_reported_as_success(self): for output in [b"not json", b'{"ok":false,"error":"failed"}']: proc = AsyncMock() proc.returncode = 0 proc.communicate.return_value = (output, b'') with patch("awg_gus.core.asyncio.create_subprocess_exec", return_value=proc): with self.assertRaises(AwgError): await Awg("awg2").call("version") if __name__ == "__main__": unittest.main()