291 lines
13 KiB
Python
291 lines
13 KiB
Python
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()
|