hz4th_coder
c877b25ff6
- collector.py: 机器A端 FastAPI 服务(16018)，token认证 + strict/open 命令白名单 + SQLite存储 - agent.sh: 机器B端轻量agent（仅bash+curl+base64，零安装），采集CPU/内存/磁盘/负载/开机时间 + 长轮询执行命令回传结果 - hostctl.py: 机器A端 CLI（status/hosts/run/history/commands） - host-agent.service: 机器B端 systemd 服务 - start.sh: collector 启停脚本
404 lines
14 KiB
Python
404 lines
14 KiB
Python
#!/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")
|