#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
remote-host-agent - 机器B状态监控/控制 Collector 服务（运行在机器A，固定IP端）
方案A：HTTP 上报 + 命令队列长轮询
- POST /api/report   B 端定时上报 CPU/内存/磁盘等指标
- GET  /api/poll     B 端长轮询取命令（阻塞最多 65s）
- POST /api/result   B 端回传命令执行结果（base64）
- GET  /api/status   查询某台主机实时状态
- GET  /api/hosts    列出所有已注册主机
- GET  /api/history  查询历史指标
- POST /api/command  下发命令到某台主机的队列
- GET  /api/health   健康检查
- GET  /api/config   查看当前配置（不含 token）
"""
import os
import sys
import json
import time
import base64
import secrets
import sqlite3
import threading
import logging
from pathlib import Path
from contextlib import contextmanager

from fastapi import FastAPI, Request, HTTPException, Depends, Query
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import uvicorn

# ---------- 基础配置 ----------
BASE_DIR = Path(__file__).resolve().parent
DATA_DIR = BASE_DIR / "data"
LOG_DIR = BASE_DIR / "logs"
DATA_DIR.mkdir(exist_ok=True)
LOG_DIR.mkdir(exist_ok=True)

DB_PATH = DATA_DIR / "host_agent.db"
TOKEN_FILE = DATA_DIR / "token.txt"

HOST = os.environ.get("HOST_AGENT_HOST", "0.0.0.0")
PORT = int(os.environ.get("HOST_AGENT_PORT", "16018"))

# 命令控制模式：strict = 白名单（默认，安全）；open = 允许任意命令（仅审计）
CMD_MODE = os.environ.get("HOST_AGENT_CMD_MODE", "strict").lower()

# B 端上报心跳超时（秒），超过则标记离线
OFFLINE_AFTER = int(os.environ.get("HOST_AGENT_OFFLINE_AFTER", "180"))

# 长轮询最长阻塞时间（秒）
POLL_WAIT = 65

# ---------- 命令白名单（strict 模式） ----------
# 只允许以这些命令名开头的命令（按空格分隔的第一个 token 精确匹配）
CMD_WHITELIST = {
    "df", "free", "top", "ps", "uptime", "uname", "hostname", "whoami",
    "cat", "ls", "find", "du", "netstat", "ss", "sar", "vmstat", "iostat",
    "mpstat", "pgrep", "lsof", "date", "echo", "ip", "ifconfig", "mount",
    "sysctl", "lscpu", "lsblk", "nproc", "getconf", "id", "dmesg", "lspci",
    "pwd", "head", "tail", "grep", "wc", "file", "stat", "readlink",
    "hostnamectl", "systemctl", "journalctl", "nvidia-smi", "psutil",
}

# ---------- 日志 ----------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler(LOG_DIR / "collector.log", encoding="utf-8"),
    ],
)
log = logging.getLogger("collector")

# ---------- Token ----------
def load_or_create_token():
    if TOKEN_FILE.exists():
        token = TOKEN_FILE.read_text().strip()
        if token:
            return token
    token = os.environ.get("HOST_AGENT_TOKEN") or secrets.token_urlsafe(32)
    TOKEN_FILE.write_text(token)
    log.info("已%s token 文件: %s", "更新" if os.environ.get("HOST_AGENT_TOKEN") else "生成", TOKEN_FILE)
    return token

TOKEN = load_or_create_token()

# ---------- 数据库 ----------
def get_conn():
    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
    conn.row_factory = sqlite3.Row
    return conn

def init_db():
    with get_conn() as conn:
        conn.executescript("""
        CREATE TABLE IF NOT EXISTS hosts (
            host        TEXT PRIMARY KEY,
            name        TEXT,
            first_seen  REAL,
            last_seen   REAL,
            online      INTEGER DEFAULT 0,
            latest      TEXT
        );
        CREATE TABLE IF NOT EXISTS metrics (
            id   INTEGER PRIMARY KEY AUTOINCREMENT,
            host TEXT,
            ts   REAL,
            cpu  REAL,
            mem  REAL,
            disk REAL,
            load TEXT,
            uptime INTEGER,
            extra TEXT
        );
        CREATE INDEX IF NOT EXISTS idx_metrics_host_ts ON metrics(host, ts);
        CREATE TABLE IF NOT EXISTS commands (
            id        INTEGER PRIMARY KEY AUTOINCREMENT,
            host      TEXT,
            cmd       TEXT,
            timeout   INTEGER DEFAULT 30,
            status    TEXT DEFAULT 'pending',  -- pending/sent/done/failed
            note      TEXT,
            created_at REAL,
            sent_at   REAL,
            result    TEXT,
            result_at REAL
        );
        CREATE INDEX IF NOT EXISTS idx_commands_host_status ON commands(host, status);
        """)

# ---------- 命令队列唤醒（长轮询用） ----------
_host_events = {}
_host_events_lock = threading.Lock()

def event_for(host):
    with _host_events_lock:
        if host not in _host_events:
            _host_events[host] = threading.Event()
        return _host_events[host]

def notify_host(host):
    event_for(host).set()

def reset_host_event(host):
    event_for(host).clear()

# ---------- 认证 ----------
def check_auth(authorization: str):
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="缺少 Bearer token")
    if authorization[len("Bearer "):].strip() != TOKEN:
        raise HTTPException(status_code=401, detail="token 无效")

# ---------- 命令校验 ----------
def validate_command(cmd: str) -> str:
    cmd = cmd.strip()
    if not cmd:
        raise HTTPException(status_code=400, detail="命令不能为空")
    if len(cmd) > 2048:
        raise HTTPException(status_code=400, detail="命令过长")
    if CMD_MODE == "strict":
        first = cmd.split()[0].lstrip("$").strip()
        if first not in CMD_WHITELIST:
            raise HTTPException(
                status_code=403,
                detail=f"strict 模式下命令被白名单拦截: {first}。可用 HOST_AGENT_CMD_MODE=open 放开",
            )
    return cmd

# ---------- FastAPI ----------
app = FastAPI(title="Remote Host Agent", version="1.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.on_event("startup")
def on_startup():
    init_db()
    log.info("Collector 启动: %s:%s  CMD_MODE=%s", HOST, PORT, CMD_MODE)
    log.info("Token 已写入 %s（agent.sh / hostctl 需使用）", TOKEN_FILE)

@app.get("/api/health")
def health():
    return {"status": "ok", "time": time.time()}

@app.get("/api/config")
def config():
    return {
        "port": PORT,
        "cmd_mode": CMD_MODE,
        "offline_after": OFFLINE_AFTER,
        "whitelist_count": len(CMD_WHITELIST) if CMD_MODE == "strict" else None,
    }

# ---------- B 端上报 ----------
@app.post("/api/report")
def report(
    request: Request,
    host: str = Query(..., description="主机标识"),
    cpu: float = Query(None),
    mem: float = Query(None),
    disk: float = Query(None),
    load: str = Query(None),
    uptime: int = Query(None),
):
    check_auth(request.headers.get("authorization", ""))
    host = host.strip()[:64]
    now = time.time()
    extra = {}
    for k, v in request.query_params.items():
        if k.startswith("x_"):
            extra[k[2:]] = v

    with get_conn() as conn:
        conn.execute(
            "INSERT INTO metrics(host,ts,cpu,mem,disk,load,uptime,extra) VALUES(?,?,?,?,?,?,?,?)",
            (host, now, cpu, mem, disk, load, uptime, json.dumps(extra, ensure_ascii=False)),
        )
        conn.execute("""
            INSERT INTO hosts(host,name,first_seen,last_seen,online,latest)
            VALUES(?,?,?,?,1,?)
            ON CONFLICT(host) DO UPDATE SET
              last_seen=excluded.last_seen, online=1, latest=excluded.latest
        """, (host, host, now, now, json.dumps({
            "cpu": cpu, "mem": mem, "disk": disk, "load": load,
            "uptime": uptime, "extra": extra, "ts": now,
        }, ensure_ascii=False)))
    return {"ok": True, "time": now}

# ---------- B 端长轮询取命令 ----------
@app.get("/api/poll")
def poll(request: Request, host: str = Query(...)):
    check_auth(request.headers.get("authorization", ""))
    host = host.strip()[:64]
    # 先标记在线
    with get_conn() as conn:
        conn.execute("UPDATE hosts SET last_seen=?, online=1 WHERE host=?", (time.time(), host))
    deadline = time.time() + POLL_WAIT
    ev = event_for(host)
    while time.time() < deadline:
        with get_conn() as conn:
            row = conn.execute(
                "SELECT id,host,cmd,timeout FROM commands WHERE host=? AND status='pending' ORDER BY id LIMIT 1",
                (host,),
            ).fetchone()
        if row:
            with get_conn() as conn:
                conn.execute(
                    "UPDATE commands SET status='sent', sent_at=? WHERE id=?",
                    (time.time(), row["id"]),
                )
            reset_host_event(host)
            return {
                "cmd_id": row["id"],
                "cmd": row["cmd"],
                "timeout": row["timeout"],
            }
        # 等待新的命令（被 POST /api/command 唤醒）
        ev.wait(timeout=min(2.0, deadline - time.time()))
    reset_host_event(host)
    return {"cmd_id": None, "cmd": None}

# ---------- B 端回传结果 ----------
@app.post("/api/result")
def result(
    request: Request,
    host: str = Query(...),
    cmd_id: int = Query(...),
    output: str = Query(..., description="命令输出（base64 编码）"),
    exit_code: int = Query(0),
):
    check_auth(request.headers.get("authorization", ""))
    host = host.strip()[:64]
    try:
        out_text = base64.b64decode(output).decode("utf-8", errors="replace")
    except Exception:
        out_text = "[解码失败]"
    with get_conn() as conn:
        cur = conn.execute(
            "UPDATE commands SET status=?, result=?, result_at=? WHERE id=? AND host=?",
            ("done" if exit_code == 0 else "failed", out_text, time.time(), cmd_id, host),
        )
    if cur.rowcount == 0:
        return {"ok": False, "detail": "未找到对应命令"}
    log.info("命令 #%s (%s) 完成 exit=%s", cmd_id, host, exit_code)
    return {"ok": True}

# ---------- 查询 ----------
@app.get("/api/hosts")
def hosts(request: Request):
    check_auth(request.headers.get("authorization", ""))
    now = time.time()
    with get_conn() as conn:
        rows = conn.execute("SELECT * FROM hosts").fetchall()
    result = []
    for r in rows:
        latest = json.loads(r["latest"]) if r["latest"] else {}
        online = (now - r["last_seen"]) < OFFLINE_AFTER
        result.append({
            "host": r["host"],
            "name": r["name"],
            "online": online,
            "last_seen": r["last_seen"],
            "first_seen": r["first_seen"],
            "latest": latest,
        })
    # 后台清理离线标记（延迟处理，不影响返回）
    if result:
        with get_conn() as conn:
            conn.execute("UPDATE hosts SET online=0 WHERE ? - last_seen > ?", (now, OFFLINE_AFTER))
    return {"hosts": result}

@app.get("/api/status")
def status(request: Request, host: str = Query(...)):
    check_auth(request.headers.get("authorization", ""))
    now = time.time()
    with get_conn() as conn:
        row = conn.execute("SELECT * FROM hosts WHERE host=?", (host,)).fetchone()
        if not row:
            raise HTTPException(status_code=404, detail=f"主机 {host} 未注册")
        latest = json.loads(row["latest"]) if row["latest"] else {}
        online = (now - row["last_seen"]) < OFFLINE_AFTER
        last_cmd = conn.execute(
            "SELECT id,cmd,status,result,created_at,result_at FROM commands "
            "WHERE host=? AND status IN ('done','failed') ORDER BY id DESC LIMIT 1",
            (host,),
        ).fetchone()
    data = {
        "host": row["host"],
        "name": row["name"],
        "online": online,
        "last_seen": row["last_seen"],
        "age_sec": round(now - row["last_seen"], 1),
        "latest": latest,
    }
    if last_cmd:
        data["last_command"] = dict(last_cmd)
    return data

@app.get("/api/history")
def history(
    request: Request,
    host: str = Query(...),
    limit: int = Query(30, ge=1, le=1000),
):
    check_auth(request.headers.get("authorization", ""))
    with get_conn() as conn:
        rows = conn.execute(
            "SELECT ts,cpu,mem,disk,load,uptime FROM metrics WHERE host=? ORDER BY ts DESC LIMIT ?",
            (host, limit),
        ).fetchall()
    return {"host": host, "count": len(rows), "points": [dict(r) for r in rows]}

@app.get("/api/commands")
def list_commands(
    request: Request,
    host: str = Query(None),
    limit: int = Query(30, ge=1, le=200),
):
    check_auth(request.headers.get("authorization", ""))
    with get_conn() as conn:
        if host:
            rows = conn.execute(
                "SELECT * FROM commands WHERE host=? ORDER BY id DESC LIMIT ?", (host, limit),
            ).fetchall()
        else:
            rows = conn.execute("SELECT * FROM commands ORDER BY id DESC LIMIT ?", (limit,)).fetchall()
    return {"commands": [dict(r) for r in rows]}

# ---------- 下发命令 ----------
@app.post("/api/command")
def command(
    request: Request,
    payload: dict,
):
    check_auth(request.headers.get("authorization", ""))
    host = str(payload.get("host", "")).strip()[:64]
    cmd = str(payload.get("cmd", "")).strip()
    timeout = int(payload.get("timeout", 30))
    note = str(payload.get("note", ""))[:200]
    if not host:
        raise HTTPException(status_code=400, detail="缺少 host")
    cmd = validate_command(cmd)
    timeout = min(max(timeout, 1), 300)
    now = time.time()
    with get_conn() as conn:
        cur = conn.execute(
            "INSERT INTO commands(host,cmd,timeout,status,note,created_at) VALUES(?,?,?,?,?,?)",
            (host, cmd, timeout, "pending", note, now),
        )
        cmd_id = cur.lastrowid
    notify_host(host)
    log.info("下发命令 #%s -> %s: %s", cmd_id, host, cmd)
    return {"ok": True, "cmd_id": cmd_id}

if __name__ == "__main__":
    uvicorn.run(app, host=HOST, port=PORT, log_level="info")