v2.5.0: align key controls and add IP limits

admin-web/server.py

@@ -53,6 +53,7 @@ USER_RE = re.compile(r"^[A-Za-z0-9_.-]{1,48}$")
 LANG_RE = re.compile(r"^(en|ru)$")
 SENSITIVE_CONFIG_KEYS = {"secret"}
 BACKUP_NAME_RE = re.compile(r"^[A-Za-z0-9_.-]+\.tar\.gz(\.enc)?$")
+MAX_UNIQUE_IP_LIMIT = 1000000
 TRAFFIC_WINDOWS = {
     "15m": 15 * 60,
     "1h": 60 * 60,
@@ -197,6 +198,38 @@ def read_telemt_users() -> dict[str, str]:
     return users
 
 
+def read_toml_int_table(table: str) -> dict[str, int]:
+    if not TELEMT_CONFIG.exists():
+        return {}
+    values: dict[str, int] = {}
+    section = f"[{table}]"
+    in_table = False
+    for raw in TELEMT_CONFIG.read_text(encoding="utf-8", errors="ignore").splitlines():
+        line = raw.strip()
+        if line == section:
+            in_table = True
+            continue
+        if in_table and line.startswith("["):
+            break
+        if not in_table or not line or line.startswith("#") or "=" not in line:
+            continue
+        name, value = line.split("=", 1)
+        name = parse_toml_key(name)
+        if not USER_RE.match(name):
+            continue
+        raw_value = value.strip().split("#", 1)[0].strip().strip('"').strip("'")
+        try:
+            number = int(raw_value)
+        except ValueError:
+            continue
+        values[name] = max(0, number)
+    return values
+
+
+def read_user_max_unique_ips() -> dict[str, int]:
+    return read_toml_int_table("access.user_max_unique_ips")
+
+
 def read_disabled_users() -> dict[str, str]:
     raw = load_json(DISABLED_USERS_FILE, {}) or {}
     if not isinstance(raw, dict):
@@ -227,11 +260,12 @@ def write_disabled_users(users: dict[str, str]) -> None:
 def read_user_records() -> dict[str, dict[str, Any]]:
     active = read_telemt_users()
     disabled = read_disabled_users()
+    ip_limits = read_user_max_unique_ips()
     records: dict[str, dict[str, Any]] = {}
     for name, secret in disabled.items():
-        records[name] = {"secret": secret, "enabled": False}
+        records[name] = {"secret": secret, "enabled": False, "max_unique_ips": ip_limits.get(name, 0)}
     for name, secret in active.items():
-        records[name] = {"secret": secret, "enabled": True}
+        records[name] = {"secret": secret, "enabled": True, "max_unique_ips": ip_limits.get(name, 0)}
     return records
 
 
@@ -243,6 +277,25 @@ def _ordered_user_lines(users: dict[str, str]) -> list[str]:
     return [f'{quote_toml_key(name)} = "{users[name]}"' for name in names]
 
 
+def _ordered_user_int_lines(values: dict[str, int]) -> list[str]:
+    positive: dict[str, int] = {}
+    for name, value in values.items():
+        name_s = str(name)
+        if not USER_RE.match(name_s):
+            continue
+        try:
+            number = int(value)
+        except (TypeError, ValueError):
+            continue
+        if number > 0:
+            positive[name_s] = number
+    names = []
+    if "main" in positive:
+        names.append("main")
+    names.extend(sorted(n for n in positive if n != "main"))
+    return [f'{quote_toml_key(name)} = {positive[name]}' for name in names]
+
+
 def parse_toml_key(raw: str) -> str:
     key = raw.strip()
     if len(key) >= 2 and key[0] == key[-1] == '"':
@@ -293,6 +346,55 @@ def write_telemt_users(users: dict[str, str]) -> None:
     tmp.replace(TELEMT_CONFIG)
 
 
+def write_toml_int_table(table: str, values: dict[str, int]) -> None:
+    TELEMT_CONFIG.parent.mkdir(parents=True, exist_ok=True)
+    lines = TELEMT_CONFIG.read_text(encoding="utf-8", errors="ignore").splitlines() if TELEMT_CONFIG.exists() else []
+    rendered = _ordered_user_int_lines(values)
+    header = f"[{table}]"
+    out: list[str] = []
+    in_table = False
+    found = False
+
+    for raw in lines:
+        if raw.strip() == header:
+            found = True
+            in_table = True
+            if rendered:
+                out.append(raw)
+                out.extend(rendered)
+            continue
+        if in_table and raw.strip().startswith("["):
+            in_table = False
+        if in_table:
+            continue
+        out.append(raw)
+
+    if not found and rendered:
+        if out and out[-1].strip():
+            out.append("")
+        out.append(header)
+        out.extend(rendered)
+
+    tmp = TELEMT_CONFIG.with_name(TELEMT_CONFIG.name + ".tmp")
+    tmp.write_text("\n".join(out).rstrip() + "\n", encoding="utf-8")
+    os.chmod(tmp, 0o600)
+    tmp.replace(TELEMT_CONFIG)
+
+
+def write_user_max_unique_ips(values: dict[str, int]) -> None:
+    write_toml_int_table("access.user_max_unique_ips", values)
+
+
+def normalize_max_unique_ips(value: Any) -> int:
+    try:
+        number = int(value)
+    except (TypeError, ValueError):
+        raise ValueError("max_unique_ips must be an integer") from None
+    if number < 0 or number > MAX_UNIQUE_IP_LIMIT:
+        raise ValueError(f"max_unique_ips must be between 0 and {MAX_UNIQUE_IP_LIMIT}")
+    return number
+
+
 def restart_service(name: str) -> bool:
     code, _, _ = run(["systemctl", "restart", name], timeout=25)
     if code != 0:
@@ -1071,6 +1173,7 @@ def user_payload(
     name: str,
     secret: str,
     enabled: bool = True,
+    max_unique_ips: int = 0,
     include_runtime: bool = False,
     traffic_snapshot: dict[str, Any] | None = None,
 ) -> dict[str, Any]:
@@ -1080,6 +1183,7 @@ def user_payload(
         "link": proxy_link(secret),
         "main": name == "main",
         "enabled": bool(enabled),
+        "max_unique_ips": _int_value(max_unique_ips),
     }
     if traffic_snapshot:
         item["traffic"] = {
@@ -1177,7 +1281,13 @@ class AdminHandler(BaseHTTPRequestHandler):
             items = []
             for name in sorted(users, key=lambda item: (item != "main", item)):
                 record = users[name]
-                items.append(user_payload(name, record["secret"], record["enabled"], traffic_snapshot=latest.get(name)))
+                items.append(user_payload(
+                    name,
+                    record["secret"],
+                    record["enabled"],
+                    record.get("max_unique_ips", 0),
+                    traffic_snapshot=latest.get(name),
+                ))
             self.send_json({"ok": True, "data": items})
         elif path.startswith("/api/users/") and path.endswith("/qr"):
             name = urllib.parse.unquote(path[len("/api/users/"):-len("/qr")])
@@ -1231,7 +1341,14 @@ class AdminHandler(BaseHTTPRequestHandler):
                 self.send_error_json(404, "user not found")
                 return
             record = users[name]
-            self.send_json({"ok": True, "data": user_payload(name, record["secret"], record["enabled"], include_runtime=True, traffic_snapshot=latest_user_stats().get(name))})
+            self.send_json({"ok": True, "data": user_payload(
+                name,
+                record["secret"],
+                record["enabled"],
+                record.get("max_unique_ips", 0),
+                include_runtime=True,
+                traffic_snapshot=latest_user_stats().get(name),
+            )})
         elif path == "/api/backups":
             self.send_json({"ok": True, "data": list_backups()})
         elif path == "/api/backups/schedule":
@@ -1296,7 +1413,38 @@ class AdminHandler(BaseHTTPRequestHandler):
                 self.send_error_json(500, f"failed to save config: {exc}")
                 return
             restart_requested = request_service_restart("telemt")
-            self.send_json({"ok": True, "data": user_payload(name, secret, True), "restart": {"mode": "async", "requested": restart_requested}})
+            self.send_json({"ok": True, "data": user_payload(name, secret, True, 0), "restart": {"mode": "async", "requested": restart_requested}})
+        elif path.startswith("/api/users/") and path.endswith("/max-ips"):
+            name = urllib.parse.unquote(path[len("/api/users/"):-len("/max-ips")])
+            try:
+                limit = normalize_max_unique_ips(body.get("max_unique_ips"))
+            except ValueError as exc:
+                self.send_error_json(400, str(exc))
+                return
+            try:
+                with FileLock(USER_LOCK_FILE):
+                    records = read_user_records()
+                    if name not in records:
+                        self.send_error_json(404, "user not found")
+                        return
+                    limits = read_user_max_unique_ips()
+                    if limit > 0:
+                        limits[name] = limit
+                    else:
+                        limits.pop(name, None)
+                    write_user_max_unique_ips(limits)
+                    record = read_user_records()[name]
+            except Exception as exc:
+                self.send_error_json(500, f"failed to save config: {exc}")
+                return
+            restart_requested = request_service_restart("telemt")
+            self.send_json({"ok": True, "data": user_payload(
+                name,
+                record["secret"],
+                record["enabled"],
+                record.get("max_unique_ips", 0),
+                traffic_snapshot=latest_user_stats().get(name),
+            ), "restart": {"mode": "async", "requested": restart_requested}})
         elif path.startswith("/api/users/") and path.endswith("/enabled"):
             name = urllib.parse.unquote(path[len("/api/users/"):-len("/enabled")])
             if name == "main":
@@ -1327,7 +1475,7 @@ class AdminHandler(BaseHTTPRequestHandler):
                 self.send_error_json(500, f"failed to save config: {exc}")
                 return
             restart_requested = request_service_restart("telemt")
-            self.send_json({"ok": True, "data": user_payload(name, secret, enabled), "restart": {"mode": "async", "requested": restart_requested}})
+            self.send_json({"ok": True, "data": user_payload(name, secret, enabled, records[name].get("max_unique_ips", 0)), "restart": {"mode": "async", "requested": restart_requested}})
         elif path == "/api/backups":
             ok, result = create_backup()
             self.send_json({"ok": ok, "data": {"path": result, "backups": list_backups()}}, 200 if ok else 500)
@@ -1399,8 +1547,11 @@ class AdminHandler(BaseHTTPRequestHandler):
                     return
                 active.pop(name, None)
                 disabled.pop(name, None)
+                limits = read_user_max_unique_ips()
+                limits.pop(name, None)
                 write_telemt_users(active)
                 write_disabled_users(disabled)
+                write_user_max_unique_ips(limits)
         except Exception as exc:
             self.send_error_json(500, f"failed to save config: {exc}")
             return