#!/usr/bin/env python3
"""Runs ON the box. Makes a DB-provisioned node upgrade-safe.

WHY THIS EXISTS
---------------
3x-ui's real source of truth for client credentials is the per-inbound
`inbounds.settings.clients` JSON array -- NOT the normalized `clients` /
`client_inbounds` tables. Those tables are a *derived projection* of the JSON:

  * client_link.go SyncInbound(tx, inboundId, clients) does an UNCONDITIONAL
    `DELETE FROM client_inbounds WHERE inbound_id = ?` and then rebuilds the
    rows only from the `clients` slice it was handed. Empty slice -> delete,
    insert nothing.
  * inbound_migration.go MigrationRequirements() hands SyncInbound the clients
    it got from GetClients(inbound), i.e. parsed from the settings JSON.
  * The panel UI path (client_inbound_apply.go addInboundClient) ALWAYS writes
    the client into the settings JSON first, saves the inbound, then projects
    to the tables via SyncInbound. So a UI-managed install always has a
    populated settings.clients and migration is a harmless no-op for it.

A node provisioned by writing the tables directly (settings.clients left as
`[]`) is therefore off-model: any run of the migration path -- `x-ui migrate`,
which install.sh/update.sh invoke on upgrade, plus DB import/restore and the
SQLite->PostgreSQL copy -- wipes every client_inbounds row and silently drops
all credentials from the generated xray config. Verified empirically: running
`x-ui migrate` on a table-only node emptied client_inbounds immediately.
This is NOT specific to any one release; the logic is unchanged since 3.5.0.

WHAT THIS DOES
--------------
Backfills `settings.clients` for every proxy inbound from the node-main client
row, then re-asserts the client_inbounds links. Idempotent -- safe to re-run,
and re-running after any 3x-ui upgrade is the recommended belt-and-braces step.

FIELD SEMANTICS (client_link.go applyClientRecordMerge) -- get these wrong and
you corrupt the merged client row:
  * Credential-ish fields overwrite ONLY when non-empty: id, password, auth,
    secret, security, privateKey, publicKey. That's why a vmess entry may carry
    only `id` and a trojan/ss entry only `password` -- across the six inbounds
    the same email merges to a union instead of clobbering.
  * These overwrite UNCONDITIONALLY, so all six entries must agree:
    enable, subId, limitIp, totalGB, expiryTime, tgId, comment, reset, flow.
  * `enable` is the dangerous one: omit it and JSON unmarshals it to false,
    which makes xray.go skip the client on EVERY inbound -> node dead.
"""
import json, os, sqlite3, subprocess, sys, time

DB = os.environ.get("XUI_DB", "/etc/x-ui/x-ui.db")
EMAIL = os.environ.get("XUI_EMAIL", "node-main")
RESTART = os.environ.get("RESTART_XUI", "1") == "1"

PROXY_PROTOCOLS = ("vmess", "vless", "trojan", "shadowsocks")


def build_entry(protocol, uuid, password, sub_id, method, now):
    """One settings.clients entry. Unconditional fields identical across inbounds."""
    e = {
        "email": EMAIL,
        "enable": True,          # MUST be explicit -- see module docstring
        "subId": sub_id or "",
        "flow": "",
        "limitIp": 0,
        "totalGB": 0,
        "expiryTime": 0,
        "tgId": 0,
        "comment": "",
        "reset": 0,
        "created_at": now,
        "updated_at": now,
    }
    if protocol in ("vmess", "vless"):
        e["id"] = uuid
        e["security"] = "auto"
        if protocol == "vmess":
            e["alterId"] = 0
    else:
        e["password"] = password
        if protocol == "shadowsocks":
            e["method"] = method
    return e


def main():
    con = sqlite3.connect(DB)
    cur = con.cursor()

    row = cur.execute(
        "SELECT id, uuid, password, sub_id FROM clients WHERE email=?", (EMAIL,)
    ).fetchone()
    if not row:
        sys.exit("ERROR: no client row for email=%r -- nothing to embed" % EMAIL)
    cid, uuid, password, sub_id = row
    now = int(time.time() * 1000)

    inbounds = cur.execute(
        "SELECT id, port, protocol, settings FROM inbounds ORDER BY port"
    ).fetchall()

    changed = 0
    for iid, port, protocol, settings_raw in inbounds:
        if protocol not in PROXY_PROTOCOLS:
            continue
        try:
            st = json.loads(settings_raw) if settings_raw else {}
        except (TypeError, ValueError):
            print("  :%s %s -- settings not valid JSON, SKIPPED" % (port, protocol))
            continue

        method = st.get("method", "chacha20-poly1305")
        entry = build_entry(protocol, uuid, password, sub_id, method, now)

        existing = st.get("clients")
        if not isinstance(existing, list):
            existing = []
        # Replace only our own email's entry; leave any other client untouched.
        kept = [c for c in existing if isinstance(c, dict) and c.get("email") != EMAIL]
        new_clients = kept + [entry]

        if existing == new_clients:
            print("  :%s %-12s settings.clients already correct" % (port, protocol))
        else:
            st["clients"] = new_clients
            cur.execute(
                "UPDATE inbounds SET settings=? WHERE id=?",
                (json.dumps(st, indent=2, ensure_ascii=False), iid),
            )
            changed += 1
            print("  :%s %-12s settings.clients embedded" % (port, protocol))

        # Re-assert the derived link (harmless if already there).
        cur.execute(
            "INSERT OR IGNORE INTO client_inbounds (client_id,inbound_id,flow_override,created_at) "
            "VALUES (?,?,?,?)", (cid, iid, "", now),
        )

    # client_traffics.email is globally unique -> exactly one row for this client.
    if not cur.execute("SELECT 1 FROM client_traffics WHERE email=?", (EMAIL,)).fetchone():
        first = cur.execute("SELECT id FROM inbounds ORDER BY port").fetchone()
        if first:
            cur.execute(
                "INSERT INTO client_traffics (inbound_id,enable,email,up,down,expiry_time,total,reset) "
                "VALUES (?,1,?,0,0,0,0,0)", (first[0], EMAIL),
            )

    con.commit()
    links = cur.execute("SELECT count(*) FROM client_inbounds").fetchone()[0]
    con.close()
    print("  settings.clients updated on %d inbound(s); client_inbounds rows now: %d" % (changed, links))

    if RESTART:
        subprocess.run(["systemctl", "restart", "x-ui"], check=True)
        time.sleep(6)
        print("OK -- x-ui restarted; node is now upgrade-safe (settings.clients is authoritative)")


if __name__ == "__main__":
    main()