"""Local AWG API, private state, and GusVoice delivery. No shell commands from chat.""" from __future__ import annotations import asyncio import contextlib import html import json import logging import os from pathlib import Path import re import signal import sqlite3 import time from gusvoice_bot import AsyncClient, GusVoiceApiError log = logging.getLogger("awg-gus") NAME = re.compile(r"[A-Za-z0-9_-]{1,32}\Z") LIMIT = re.compile(r"(?:[1-9][0-9]{0,5}[MG]|off)\Z") SPEED = re.compile(r"(?:[1-9][0-9]{0,5}|off)\Z") EXPIRE = re.compile(r"\+[1-9][0-9]{0,3}[dh]\Z") def esc(value): return html.escape(str(value if value is not None else "")) def redacted(value): text = str(value) text = re.sub(r"(?im)(PrivateKey|PresharedKey|BOT_TOKEN|password)\s*[:=]\s*\S+", r"\1=[скрыто]", text) return re.sub(r"gvb_[A-Za-z0-9_-]+", "[токен скрыт]", text) def size(value): n = float(value or 0) for unit in ("Б", "КиБ", "МиБ", "ГиБ", "ТиБ"): if abs(n) < 1024 or unit == "ТиБ": return f"{n:.1f} {unit}" n /= 1024 class AwgError(RuntimeError): pass def validate_command(args): """Second allow-list immediately before a privileged subprocess starts.""" a = tuple(args) if a in {("version",), ("status",), ("server", "info"), ("server", "restart"), ("clients", "list"), ("clients", "export"), ("diag", "status"), ("backup", "create"), ("backup", "list"), ("tunnels", "status")}: return if len(a) >= 3 and a[0] == "client" and NAME.fullmatch(a[2]): op = a[1] if len(a) == 3 and op in {"conf", "del", "reissue", "unexpire", "limit-reset"}: return if len(a) == 4: if op == "rename" and NAME.fullmatch(a[3]): return if op == "expire" and EXPIRE.fullmatch(a[3]): return if op == "speed" and SPEED.fullmatch(a[3]): return if len(a) == 5 and op == "limit" and LIMIT.fullmatch(a[3]) and a[4] in {"month", "total"}: return if op == "add" and len(a) in {4, 5} and a[3] == "mimicry=server": if len(a) == 4 or (a[4].startswith("expire=") and EXPIRE.fullmatch(a[4][7:])): return if len(a) == 4 and a[:2] == ("traffic", "daily") and (a[2] == "all" or NAME.fullmatch(a[2])) and a[3] == "7": return raise ValueError("Недопустимая операция awg2") class Awg: def __init__(self, binary="/usr/local/bin/awg2"): self.binary = binary self.lock = asyncio.Lock() async def call(self, *args, timeout=120): validate_command(args) # Upstream awg2 also uses a flock shared with the Telegram bot. async with self.lock: try: proc = await asyncio.create_subprocess_exec( self.binary, "api", *args, stdin=asyncio.subprocess.DEVNULL, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, start_new_session=(os.name == "posix")) except OSError as e: raise AwgError("Не удалось запустить awg2. Проверьте установку и путь к программе.") from e try: out, _ = await asyncio.wait_for(proc.communicate(), timeout) except (asyncio.TimeoutError, asyncio.CancelledError) as e: with contextlib.suppress(ProcessLookupError): if os.name == "posix": os.killpg(proc.pid, signal.SIGTERM) else: proc.kill() try: await asyncio.wait_for(proc.communicate(), 5) except asyncio.TimeoutError: with contextlib.suppress(ProcessLookupError): if os.name == "posix": os.killpg(proc.pid, signal.SIGKILL) else: proc.kill() await proc.communicate() if isinstance(e, asyncio.CancelledError): raise raise AwgError("awg2 не ответил вовремя. Проверьте результат в меню: операция могла выполниться.") from e try: result = json.loads(out.decode("utf-8", errors="replace").strip().splitlines()[-1]) if not isinstance(result, dict) or not isinstance(result.get("ok"), bool): raise ValueError except (ValueError, IndexError) as e: raise AwgError("awg2 вернул ответ без JSON API. Нужна версия AWG Toolza с командой awg2 api.") from e if not result["ok"] or proc.returncode: raise AwgError(redacted(result.get("error") or "awg2 завершился с ошибкой")[:500]) return result class State: def __init__(self, path, identity): path = Path(path) path.parent.mkdir(parents=True, exist_ok=True, mode=0o700) self.db = sqlite3.connect(path) if os.name == "posix": path.chmod(0o600) self.db.executescript(""" CREATE TABLE IF NOT EXISTS kv(key TEXT PRIMARY KEY, value TEXT NOT NULL); CREATE TABLE IF NOT EXISTS chats(uid TEXT PRIMARY KEY, chat TEXT NOT NULL); CREATE TABLE IF NOT EXISTS watches(uid TEXT, name TEXT, PRIMARY KEY(uid,name)); CREATE TABLE IF NOT EXISTS pending(id INTEGER PRIMARY KEY, uid TEXT, chat TEXT); """) old = self.get("identity") if old is not None and old != identity: self.db.close() raise ValueError("Состояние принадлежит другому боту/инстансу. Укажите отдельный state_dir.") self.set("identity", identity) def get(self, key, default=None): row = self.db.execute("SELECT value FROM kv WHERE key=?", (key,)).fetchone() return json.loads(row[0]) if row else default def set(self, key, value): with self.db: self.db.execute("INSERT OR REPLACE INTO kv VALUES (?,?)", (key, json.dumps(value))) def claim(self, update_id, uid=None, chat=None): if update_id < self.get("offset", 0): return False # Persist BEFORE any mutation. Interrupted operations are reported, never replayed. with self.db: self.db.execute("INSERT OR REPLACE INTO kv VALUES ('offset',?)", (json.dumps(update_id + 1),)) if uid and chat: self.db.execute("INSERT INTO pending VALUES (?,?,?)", (update_id, uid, chat)) self.db.execute("INSERT OR REPLACE INTO chats VALUES (?,?)", (uid, chat)) return True def finish(self, update_id): with self.db: self.db.execute("DELETE FROM pending WHERE id=?", (update_id,)) def watch(self, uid, name): with self.db: if self.watched(uid, name): self.db.execute("DELETE FROM watches WHERE uid=? AND name=?", (uid, name)) return False self.db.execute("INSERT INTO watches VALUES (?,?)", (uid, name)) return True def watched(self, uid, name): return bool(self.db.execute("SELECT 1 FROM watches WHERE uid=? AND name=?", (uid, name)).fetchone()) class Delivery: """Throttle sends/edits together; acknowledge callbacks independently.""" def __init__(self, client): self.client = client self.locks = {} self.next_send = {} async def call(self, method, **params): chat = params.get("chat_id") if chat and method in {"send_message", "edit_message_text", "edit_message_reply_markup"}: lock = self.locks.setdefault(chat, asyncio.Lock()) async with lock: await asyncio.sleep(max(0, self.next_send.get(chat, 0) - time.monotonic())) try: return await getattr(self.client, method)(**params) finally: self.next_send[chat] = time.monotonic() + 1.1 return await getattr(self.client, method)(**params) async def ack(self, callback_id, text=None): try: await self.call("answer_callback_query", callback_query_id=callback_id, text=text) except GusVoiceApiError as e: if e.code != 400: raise async def show(self, chat, text, buttons=None, message=None): params = dict(chat_id=chat, text=text, parse_mode="HTML", reply_markup={"inline_keyboard": buttons} if buttons else None) if message: try: return await self.call("edit_message_text", message_id=message, **params) except GusVoiceApiError as e: if e.code != 400: raise return await self.call("send_message", **params) def client_for(config): return AsyncClient(config["token"], config["url"], discover=False, max_retries=3)