216 lines
8.9 KiB
Python
216 lines
8.9 KiB
Python
"""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)
|