vpn/backend/db.py
2026-09-23 01:19:36 +00:00

183 lines
7.5 KiB
Python

from __future__ import annotations
import secrets
import sqlite3
import uuid
from datetime import UTC, datetime
from pathlib import Path
from backend.security import SecretBox, token_hash, token_matches
def now_iso() -> str:
return datetime.now(UTC).isoformat()
class Database:
def __init__(self, path: str, secrets_box: SecretBox) -> None:
self.path = Path(path)
self.path.parent.mkdir(parents=True, exist_ok=True)
self.box = secrets_box
self._init_schema()
def connect(self) -> sqlite3.Connection:
connection = sqlite3.connect(self.path, timeout=10, isolation_level=None)
connection.row_factory = sqlite3.Row
connection.execute("PRAGMA foreign_keys = ON")
return connection
def _init_schema(self) -> None:
with self.connect() as db:
db.executescript(
"""
CREATE TABLE IF NOT EXISTS users (
id TEXT PRIMARY KEY,
name TEXT NOT NULL,
monthly_price_rub INTEGER NOT NULL CHECK(monthly_price_rub >= 0),
payment_day INTEGER NOT NULL CHECK(payment_day BETWEEN 1 AND 31),
vless_subscription_url TEXT NOT NULL DEFAULT '',
amnezia_account_name TEXT NOT NULL DEFAULT '',
freeturn_key TEXT NOT NULL DEFAULT '',
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS server_configs (
id TEXT PRIMARY KEY,
user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
server_name TEXT NOT NULL,
amnezia_config TEXT NOT NULL,
UNIQUE(user_id, server_name)
);
CREATE TABLE IF NOT EXISTS devices (
id TEXT PRIMARY KEY,
user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
token_hash TEXT NOT NULL,
bound_user_agent TEXT,
created_at TEXT NOT NULL,
last_seen_at TEXT,
last_user_agent TEXT,
revoked_at TEXT
);
"""
)
def list_users(self) -> list[dict]:
with self.connect() as db:
users = [
dict(row) for row in db.execute("SELECT * FROM users ORDER BY created_at DESC")
]
for user in users:
user["devices"] = [
dict(row)
for row in db.execute(
"SELECT id, created_at, last_seen_at, last_user_agent, bound_user_agent, revoked_at "
"FROM devices WHERE user_id = ? ORDER BY created_at DESC",
(user["id"],),
)
]
user["server_configs"] = [
dict(row)
for row in db.execute(
"SELECT id, server_name FROM server_configs WHERE user_id = ? ORDER BY server_name",
(user["id"],),
)
]
for user in users:
user.pop("vless_subscription_url", None)
user.pop("amnezia_account_name", None)
user.pop("freeturn_key", None)
return users
def create_user(self, data: dict) -> dict:
user_id = str(uuid.uuid4())
timestamp = now_iso()
with self.connect() as db:
db.execute(
"INSERT INTO users VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
(
user_id,
data["name"],
data["monthly_price_rub"],
data["payment_day"],
self.box.seal(data.get("vless_subscription_url", "")),
self.box.seal(data.get("amnezia_account_name", "")),
self.box.seal(data.get("freeturn_key", "")),
timestamp,
timestamp,
),
)
return {"id": user_id, "name": data["name"]}
def delete_user(self, user_id: str) -> bool:
with self.connect() as db:
cursor = db.execute("DELETE FROM users WHERE id = ?", (user_id,))
return cursor.rowcount == 1
def create_device(self, user_id: str, base_url: str) -> dict:
device_id = str(uuid.uuid4())
token = secrets.token_urlsafe(32)
with self.connect() as db:
exists = db.execute("SELECT 1 FROM users WHERE id = ?", (user_id,)).fetchone()
if exists is None:
raise KeyError("user not found")
db.execute(
"INSERT INTO devices (id, user_id, token_hash, created_at) VALUES (?, ?, ?, ?)",
(device_id, user_id, token_hash(token), now_iso()),
)
return {
"device_id": device_id,
"subscription_url": f"{base_url}/d/{device_id}",
"enrollment_url": f"{base_url}/d/{device_id}?enroll={token}",
"device_token": token,
}
def delete_device(self, device_id: str) -> bool:
with self.connect() as db:
cursor = db.execute("DELETE FROM devices WHERE id = ?", (device_id,))
return cursor.rowcount == 1
def bind_and_get_subscription(
self, device_id: str, token: str, user_agent: str, user_agent_device: str
) -> dict:
now = now_iso()
with self.connect() as db:
db.execute("BEGIN IMMEDIATE")
device = db.execute("SELECT * FROM devices WHERE id = ?", (device_id,)).fetchone()
if device is None or device["revoked_at"]:
raise KeyError("device not found or revoked")
if not token_matches(token, device["token_hash"]):
raise PermissionError("invalid device credential")
if user_agent_device != device_id:
raise ValueError("device id in user-agent does not match URL")
bound = device["bound_user_agent"]
if bound is not None and bound != user_agent:
raise PermissionError("device is already bound to a different installation")
if bound is None:
db.execute(
"UPDATE devices SET bound_user_agent = ? WHERE id = ?", (user_agent, device_id)
)
db.execute(
"UPDATE devices SET last_seen_at = ?, last_user_agent = ? WHERE id = ?",
(now, user_agent, device_id),
)
user = db.execute("SELECT * FROM users WHERE id = ?", (device["user_id"],)).fetchone()
configs = db.execute(
"SELECT server_name, amnezia_config FROM server_configs WHERE user_id = ?",
(device["user_id"],),
).fetchall()
if user is None:
raise KeyError("user not found")
return {
"user_id": user["id"],
"name": user["name"],
"monthly_price_rub": user["monthly_price_rub"],
"payment_day": user["payment_day"],
"vless_subscription_url": self.box.open(user["vless_subscription_url"]),
"amnezia_account_name": self.box.open(user["amnezia_account_name"]),
"amnezia": {
row["server_name"]: self.box.open(row["amnezia_config"]) for row in configs
},
"freeturn_key": self.box.open(user["freeturn_key"]),
"device_id": device_id,
"updated_at": user["updated_at"],
}