Nexus/server/infrastructure/btpanel/ssh_bootstrap.py

"""SSH remote bootstrap: enable Baota API, token, whitelist, base_url."""

from __future__ import annotations

import base64
import json
import logging
import shlex
from typing import Any

from server.domain.models import Server
from server.infrastructure.ssh.asyncssh_pool import exec_ssh_command

logger = logging.getLogger("nexus.btpanel.ssh_bootstrap")

# Runs on sub-server as root; stdout = one JSON line
_REMOTE_SCRIPT = r"""
import fcntl, json, os, secrets, sys, subprocess, tempfile

PANEL_ROOT = "/www/server/panel"
DATA_DIR = os.path.join(PANEL_ROOT, "data")
API_PATH = os.path.join(PANEL_ROOT, "config", "api.json")
LOCK_PATH = "/tmp/nexus_bt_bootstrap.lock"

def read_text(path, default=""):
    try:
        with open(path, "r", encoding="utf-8") as f:
            return f.read().strip()
    except FileNotFoundError:
        return default

def emit(obj):
    print(json.dumps(obj, ensure_ascii=False))
    sys.exit(0)


def atomic_write_text(path, content, mode):
    directory = os.path.dirname(path) or "."
    prefix = os.path.basename(path) + ".nexus."
    fd, tmp = tempfile.mkstemp(prefix=prefix, suffix=".tmp", dir=directory, text=True)
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as f:
            f.write(content)
        os.chmod(tmp, mode)
        os.replace(tmp, path)
    except Exception:
        try:
            os.unlink(tmp)
        except FileNotFoundError:
            pass
        raise


def atomic_write_json(path, data, mode=0o600):
    directory = os.path.dirname(path) or "."
    prefix = os.path.basename(path) + ".nexus."
    fd, tmp = tempfile.mkstemp(prefix=prefix, suffix=".tmp", dir=directory, text=True)
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False)
        os.chmod(tmp, mode)
        os.replace(tmp, path)
    except Exception:
        try:
            os.unlink(tmp)
        except FileNotFoundError:
            pass
        raise


def ensure_session_cleanup_ttl():
    # Make Baota session cleanup honor session_timeout.pl instead of a hard-coded 1h.
    task_path = os.path.join(PANEL_ROOT, "task.py")
    if not os.path.isfile(task_path):
        return False
    with open(task_path, "r", encoding="utf-8") as f:
        source = f.read()
    marker = "session_cleanup_timeout = int(public.get_session_timeout() or 86400)"
    patched_condition = "                if f_time > session_cleanup_timeout:\n"
    if marker in source:
        if patched_condition not in source:
            raise RuntimeError("task.py session cleanup patch incomplete")
        return False
    anchor = (
        "            s_time = time.time()\n"
        "            f_list = os.listdir(sess_path)\n"
        "            f_num = len(f_list)\n"
    )
    replacement = (
        "            s_time = time.time()\n"
        "            session_cleanup_timeout = 86400\n"
        "            try:\n"
        "                session_cleanup_timeout = int(public.get_session_timeout() or 86400)\n"
        "            except Exception:\n"
        "                pass\n"
        "            f_list = os.listdir(sess_path)\n"
        "            f_num = len(f_list)\n"
    )
    hard_coded = "                if f_time > 3600:\n"
    anchor_index = source.find(anchor)
    if anchor_index < 0:
        raise RuntimeError("task.py session cleanup anchor not found")
    patched = source[:anchor_index] + replacement + source[anchor_index + len(anchor):]
    condition_index = patched.find(hard_coded, anchor_index + len(replacement))
    if condition_index < 0:
        raise RuntimeError("task.py session cleanup threshold not found")
    patched = patched[:condition_index] + patched_condition + patched[condition_index + len(hard_coded):]
    compile(patched, task_path, "exec")
    original_mode = os.stat(task_path).st_mode & 0o7777
    backup = task_path + ".nexus-session-ttl.bak"
    if not os.path.exists(backup):
        with open(backup, "w", encoding="utf-8") as f:
            f.write(source)
        os.chmod(backup, 0o600)
    atomic_write_text(task_path, patched, original_mode)
    return True

def main_locked():
    if len(sys.argv) < 2:
        emit({"ok": False, "error": "bad_args", "message": "missing payload"})
    try:
        payload = json.loads(sys.argv[1])
    except json.JSONDecodeError:
        emit({"ok": False, "error": "bad_args", "message": "invalid payload"})
    center_ip = (payload.get("center_ip") or "").strip()
    panel_host = (payload.get("panel_host") or "").strip()
    if not center_ip:
        emit({"ok": False, "error": "missing_center_ip", "message": "center_ip required"})
    if not os.path.isdir(PANEL_ROOT):
        emit({"ok": False, "error": "bt_not_installed", "message": "panel not installed"})
    port = read_text(os.path.join(DATA_DIR, "port.pl"), "8888") or "8888"
    admin_path = read_text(os.path.join(DATA_DIR, "admin_path.pl"), "")
    ssl = os.path.exists(os.path.join(DATA_DIR, "ssl.pl"))
    scheme = "https" if ssl else "http"
    host = panel_host or "127.0.0.1"
    base = f"{scheme}://{host}:{port}"
    # API root path does not include the panel security entrance; browser login URL is probed in ssh_login.
    api = {}
    if os.path.isfile(API_PATH):
        try:
            with open(API_PATH, "r", encoding="utf-8") as f:
                api = json.load(f)
        except Exception:
            api = {}
    actions = []
    warnings = []
    session_cleanup_ttl_ok = False
    session_cleanup_ttl_patched = False
    try:
        session_cleanup_ttl_patched = ensure_session_cleanup_ttl()
        session_cleanup_ttl_ok = True
        if session_cleanup_ttl_patched:
            actions.append("patched_session_cleanup_ttl")
    except Exception as exc:
        warnings.append(f"session_cleanup_ttl_patch_failed:{exc.__class__.__name__}")
    token = (api.get("token") or "").strip()
    if not token:
        token = secrets.token_hex(16)
        actions.append("created_token")
    if not api.get("open"):
        actions.append("enabled_api")
    api["open"] = True
    api["token"] = token
    limit = api.get("limit_addr")
    if not isinstance(limit, list):
        limit = []
    if center_ip not in limit:
        limit.append(center_ip)
        actions.append("appended_whitelist")
    api["limit_addr"] = limit
    # When config changes, write api.json and reload bt so API/open/whitelist/session patch takes effect.
    if actions:
        atomic_write_json(API_PATH, api, 0o600)
        try:
            subprocess.run(["/etc/init.d/bt", "reload"], capture_output=True, timeout=30)
        except Exception:
            pass
    emit({
        "ok": True,
        "base_url": base,
        "api_key": token,
        "actions": actions,
        "warnings": warnings,
        "session_cleanup_ttl_ok": session_cleanup_ttl_ok,
        "session_cleanup_ttl_patched": session_cleanup_ttl_patched,
        "reloaded": bool(actions),
        "verify_ssl": False,
    })

def main():
    with open(LOCK_PATH, "w") as lock_file:
        fcntl.flock(lock_file, fcntl.LOCK_EX)
        return main_locked()


if __name__ == "__main__":
    main()
"""


def build_bootstrap_command(*, center_ip: str, panel_host: str) -> str:
    payload = json.dumps({"center_ip": center_ip, "panel_host": panel_host}, ensure_ascii=False)
    script_b64 = base64.b64encode(_REMOTE_SCRIPT.encode("utf-8")).decode("ascii")
    return (
        "tmp=$(mktemp /tmp/nexus_bt_bootstrap.XXXXXX.py) && "
        "trap 'rm -f \"$tmp\"' EXIT && "
        "umask 077 && "
        f"printf %s {shlex.quote(script_b64)} | base64 -d > \"$tmp\" && "
        f"python3 \"$tmp\" {shlex.quote(payload)}; "
        "rc=$?; rm -f \"$tmp\"; exit $rc"
    )


def wrap_sudo_nopasswd(command: str, ssh_user: str | None) -> str:
    """Run remote command as root when SSH user is non-root (passwordless sudo)."""
    user = (ssh_user or "root").strip() or "root"
    if user == "root":
        return command
    return f"sudo -n bash -c {shlex.quote(command)}"


def classify_ssh_result(result: dict[str, Any]) -> tuple[str, bool]:
    """Return (error_code, permanent)."""
    status = result.get("status") or ""
    stderr_full = (result.get("stderr") or "")
    stderr = stderr_full.lower()
    stdout = (result.get("stdout") or "").lower()
    if "没有可用的 ssh 凭据" in stderr_full or "没有可用的 SSH 凭据" in stderr_full:
        return "ssh_no_credentials", True
    if status == "connection_error":
        return "ssh_connection_failed", False
    if status == "timeout":
        return "ssh_timeout", False
    # SSH 认证失败需要标记为永久失败；远端脚本写 api.json 的 PermissionError 不应误判为永久失败。
    auth_markers = (
        "permission denied (publickey)",
        "permission denied (password)",
        "permission denied, please try again",
        "authentication failed",
        "auth fail",
        "no more authentication methods",
    )
    if any(m in stderr for m in auth_markers):
        return "ssh_auth_failed", True
    if "a password is required" in stderr and "sudo" in stderr:
        return "ssh_sudo_required", True
    exit_raw = result.get("exit_code")
    exit_code = int(exit_raw) if exit_raw is not None else 1
    if status != "success" or exit_code != 0:
        if "permission denied" in stdout and "api.json" in stdout:
            return "ssh_permission_denied", False
        return "ssh_command_failed", False
    return "", False


def parse_bootstrap_stdout(stdout: str) -> dict[str, Any]:
    text = (stdout or "").strip()
    if not text:
        raise ValueError("empty bootstrap output")
    line = text.splitlines()[-1].strip()
    data = json.loads(line)
    if not isinstance(data, dict):
        raise ValueError("bootstrap output not object")
    if not data.get("ok"):
        code = str(data.get("error") or "bootstrap_failed")
        raise BootstrapRemoteError(code, str(data.get("message") or code))
    return data


class BootstrapRemoteError(Exception):
    def __init__(self, code: str, message: str, *, permanent: bool = False):
        super().__init__(message)
        self.code = code
        self.permanent = permanent


async def ssh_bootstrap_panel(server: Server, *, center_ip: str) -> dict[str, Any]:
    cmd = build_bootstrap_command(center_ip=center_ip, panel_host=(server.domain or "").strip())
    cmd = wrap_sudo_nopasswd(cmd, server.username)
    result = await exec_ssh_command(server, cmd, timeout=60, login_shell=False)
    err_code, permanent = classify_ssh_result(result)
    if err_code:
        raise BootstrapRemoteError(err_code, result.get("stderr") or err_code, permanent=permanent)
    try:
        return parse_bootstrap_stdout(result.get("stdout") or "")
    except BootstrapRemoteError:
        raise
    except (json.JSONDecodeError, ValueError) as exc:
        logger.warning("bootstrap parse failed server=%s: %s stdout=%s", server.id, exc, (result.get("stdout") or "")[:200])
        raise BootstrapRemoteError("bootstrap_parse_failed", str(exc)) from exc