#!/usr/bin/env python3
"""NAS 进程实时监控后端（纯标准库，无外部依赖；可选 nethogs 提供逐进程带宽）。

提供：
  GET  /              -> 前端页面
  GET  /api/processes -> 进程列表 JSON（含每进程网络信息 + 系统总流量 + 逐进程实时带宽）
  POST /api/kill      -> 停止进程 {"pid": int, "signal": 15|9}

所有请求需带 token（?token=xxx 或 header X-Token）。
默认 token 取自环境变量 NAS_MONITOR_TOKEN，未设置则用 "iLoveJava5"。
"""
import os
import json
import subprocess
import signal
import sys
import time
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

TOKEN = os.environ.get("NAS_MONITOR_TOKEN", "iLoveJava5")
HOST = os.environ.get("NAS_MONITOR_HOST", "127.0.0.1")
PORT = int(os.environ.get("NAS_MONITOR_PORT", "8901"))
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
SELF_PID = os.getpid()

# 受保护、不允许通过本接口终止的 PID（避免把自己或系统核心干掉）
PROTECTED_PIDS = {1, SELF_PID}


# ---------------- 网络总量采样（后台线程每秒采样 /proc/net/dev）----------------
_net_lock = threading.Lock()
_net_snap = {"rx_bytes": 0, "tx_bytes": 0, "rx_rate": 0.0, "tx_rate": 0.0, "ts": 0.0}


def _read_net_bytes():
    rx = tx = 0
    try:
        with open("/proc/net/dev") as f:
            for line in f:
                if ":" not in line:
                    continue
                name, rest = line.split(":", 1)
                if name.strip() == "lo":
                    continue
                parts = rest.split()
                if len(parts) >= 16:
                    rx += int(parts[0])
                    tx += int(parts[8])
    except Exception:
        pass
    return rx, tx


def _net_sampler():
    global _net_snap
    prev_rx, prev_tx = _read_net_bytes()
    prev_t = time.time()
    while True:
        time.sleep(1)
        rx, tx = _read_net_bytes()
        now = time.time()
        dt = now - prev_t
        rx_rate = tx_rate = 0.0
        if dt > 0:
            rx_rate = max(0.0, (rx - prev_rx)) / dt
            tx_rate = max(0.0, (tx - prev_tx)) / dt
        with _net_lock:
            _net_snap = {"rx_bytes": rx, "tx_bytes": tx,
                         "rx_rate": rx_rate, "tx_rate": tx_rate, "ts": now}
        prev_rx, prev_tx, prev_t = rx, tx, now


threading.Thread(target=_net_sampler, daemon=True).start()


def get_net_total():
    with _net_lock:
        return dict(_net_snap)


# ---------------- 各进程实时带宽（nethogs 常驻抓包，可选）----------------
# nethogs 由 Entware 提供（/opt/sbin/nethogs）。后端以 root 运行可直接抓包。
# 输出 tracemode（-t）：每次刷新以 "Refreshing:" 开头，随后每行
#   /path/to/prog/<PID>/0\t<sent KB/s>\t<recv KB/s>
NETHOGS_BIN = "/opt/sbin/nethogs"
_nh_lock = threading.Lock()
_nh_snap = {}            # pid -> (sent_kbps, recv_kbps)
_nh_proc = None


def _parse_nethogs_line(line, temp):
    line = line.rstrip("\n")
    if not line:
        return
    if line.startswith("Refreshing:"):
        with _nh_lock:
            _nh_snap.update(temp)
        temp.clear()
        return
    if "\t" not in line:
        return
    parts = line.split("\t")
    if len(parts) < 3:
        return
    try:
        sent = float(parts[-2])
        recv = float(parts[-1])
    except ValueError:
        return
    path = parts[0].strip("/")
    if path.startswith("unknown"):
        return
    segs = path.split("/")
    if len(segs) < 2:
        return
    try:
        pid = int(segs[-2])
    except ValueError:
        return
    temp[pid] = (sent, recv)


def _nethogs_loop():
    """常驻：运行 nethogs 抓包，把逐进程速率写进 _nh_snap。崩溃自动重启。"""
    global _nh_proc
    while True:
        try:
            env = dict(os.environ)
            lp = env.get("LD_LIBRARY_PATH", "")
            env["LD_LIBRARY_PATH"] = ("/opt/lib" + (":" + lp if lp else ""))
            proc = subprocess.Popen(
                [NETHOGS_BIN, "-t", "-v", "0", "-d", "2", "-a"],
                stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                env=env, text=True, bufsize=1,
            )
            _nh_proc = proc
            temp = {}
            for line in proc.stdout:
                _parse_nethogs_line(line, temp)
            proc.wait()
        except Exception:
            try:
                if _nh_proc:
                    _nh_proc.kill()
            except Exception:
                pass
        time.sleep(3)


def get_nethogs_snap():
    with _nh_lock:
        return dict(_nh_snap)


threading.Thread(target=_nethogs_loop, daemon=True).start()


# ---------------- 各进程网络连接信息（inode 关联 socket）----------------
def get_net_per_pid():
    """返回 {pid: {"conns":int,"listen":[ports],"established":int}}。"""
    entries = {}
    specs = [("tcp", "/proc/net/tcp"), ("tcp6", "/proc/net/tcp6"),
             ("udp", "/proc/net/udp"), ("udp6", "/proc/net/udp6")]
    for proto, path in specs:
        try:
            with open(path) as f:
                next(f, None)
                for line in f:
                    p = line.split()
                    if len(p) < 10:
                        continue
                    inode = p[9]
                    st = p[3]
                    local = p[1]
                    lport = 0
                    if ":" in local:
                        try:
                            lport = int(local.split(":")[1], 16)
                        except ValueError:
                            pass
                    listen = (proto.startswith("tcp") and st == "0A")
                    established = (proto.startswith("tcp") and st == "01")
                    entries[inode] = {"lport": lport, "listen": listen,
                                      "established": established}
        except Exception:
            pass

    result = {}
    try:
        for pid in os.listdir("/proc"):
            if not pid.isdigit():
                continue
            fd_dir = "/proc/%s/fd" % pid
            try:
                fds = os.listdir(fd_dir)
            except OSError:
                continue
            for fd in fds:
                try:
                    link = os.readlink("%s/%s" % (fd_dir, fd))
                except OSError:
                    continue
                if link.startswith("socket:["):
                    inode = link[len("socket:["):-1]
                    e = entries.get(inode)
                    if not e:
                        continue
                    d = result.setdefault(
                        int(pid), {"conns": 0, "listen": set(), "established": 0})
                    d["conns"] += 1
                    if e["listen"]:
                        d["listen"].add(e["lport"])
                    if e["established"]:
                        d["established"] += 1
    except Exception:
        pass

    out = {}
    for pid, d in result.items():
        out[pid] = {"conns": d["conns"], "listen": sorted(d["listen"]),
                    "established": d["established"]}
    return out


# ---------------- token 校验 ----------------
def check_token(req):
    qs = req.path.split("?", 1)
    token = ""
    if len(qs) == 2:
        for kv in qs[1].split("&"):
            if kv.startswith("token="):
                token = kv[6:]
    if not token:
        token = req.headers.get("X-Token", "")
    return token == TOKEN


# ---------------- 进程列表 ----------------
def get_processes():
    try:
        out = subprocess.check_output(
            ["ps", "-eo", "pid,user:20,pcpu,pmem,rss,etime,comm,args", "--no-headers"],
            stderr=subprocess.DEVNULL,
        ).decode("utf-8", "replace")
    except Exception as e:
        return {"error": str(e)}

    net_map = get_net_per_pid()
    nh_snap = get_nethogs_snap()
    procs = []
    for line in out.splitlines():
        if not line.strip():
            continue
        parts = line.split()
        if len(parts) < 7:
            continue
        pid, user, pcpu, pmem, rss, etime, comm = parts[:7]
        args = " ".join(parts[7:]) if len(parts) > 7 else comm
        try:
            pid_i = int(pid)
        except ValueError:
            continue
        net = net_map.get(pid_i, {"conns": 0, "listen": [], "established": 0})
        nh = nh_snap.get(pid_i)
        net["rate_sent_kbps"] = nh[0] if nh else 0.0
        net["rate_recv_kbps"] = nh[1] if nh else 0.0
        procs.append({
            "pid": pid_i,
            "user": user,
            "cpu": float(pcpu) if pcpu else 0.0,
            "mem": float(pmem) if pmem else 0.0,
            "rss_kb": int(rss) if rss.isdigit() else 0,
            "etime": etime,
            "comm": comm,
            "args": args,
            "protected": pid_i in PROTECTED_PIDS,
            "net": net,
        })
    # 按 CPU 降序
    procs.sort(key=lambda p: p["cpu"], reverse=True)
    return {
        "processes": procs,
        "count": len(procs),
        "net_total": get_net_total(),
    }


def do_kill(pid, sig):
    if not isinstance(pid, int) or pid <= 0:
        return {"ok": False, "error": "invalid pid"}
    if pid in PROTECTED_PIDS:
        return {"ok": False, "error": "该进程受保护，禁止终止（含监控自身）"}
    try:
        os.kill(pid, sig)
        return {"ok": True, "pid": pid, "signal": int(sig)}
    except ProcessLookupError:
        return {"ok": False, "error": "进程不存在（可能已退出）"}
    except PermissionError:
        return {"ok": False, "error": "权限不足，无法终止该进程（可能需要 root）"}
    except Exception as e:
        return {"ok": False, "error": str(e)}


class Handler(BaseHTTPRequestHandler):
    def log_message(self, *a):
        pass  # 静默访问日志

    def _send(self, code, obj=None, ctype="application/json; charset=utf-8"):
        if obj is None:
            body = b""
        elif ctype.startswith("application/json"):
            body = json.dumps(obj, ensure_ascii=False).encode("utf-8")
        else:
            body = obj if isinstance(obj, bytes) else obj.encode("utf-8")
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        self.wfile.write(body)

    def _token_fail(self):
        self._send(401, {"error": "token 无效"})

    def do_GET(self):
        path = self.path.split("?", 1)[0]
        if path in ("/", "/index.html"):
            try:
                with open(os.path.join(BASE_DIR, "index.html"), "rb") as f:
                    self._send(200, f.read(), "text/html; charset=utf-8")
            except FileNotFoundError:
                self._send(404, {"error": "index.html 缺失"})
            return
        if path == "/api/processes":
            if not check_token(self):
                return self._token_fail()
            self._send(200, get_processes())
            return
        if path == "/api/health":
            self._send(200, {"ok": True})
            return
        self._send(404, {"error": "not found"})

    def do_POST(self):
        path = self.path.split("?", 1)[0]
        if path == "/api/kill":
            if not check_token(self):
                return self._token_fail()
            try:
                length = int(self.headers.get("Content-Length", 0))
                raw = self.rfile.read(length) if length else b"{}"
                data = json.loads(raw.decode("utf-8", "replace") or "{}")
            except Exception:
                data = {}
            pid = data.get("pid")
            sig = data.get("signal", 15)
            try:
                sig = int(sig)
            except Exception:
                sig = 15
            if sig not in (signal.SIGTERM, signal.SIGKILL, 15, 9):
                sig = signal.SIGTERM
            if not isinstance(pid, int):
                try:
                    pid = int(pid)
                except Exception:
                    return self._send(400, {"ok": False, "error": "pid 必须是整数"})
            self._send(200, do_kill(pid, sig))
            return
        self._send(404, {"error": "not found"})


if __name__ == "__main__":
    srv = ThreadingHTTPServer((HOST, PORT), Handler)
    print(f"NAS Monitor listening on http://{HOST}:{PORT} (token={TOKEN})")
    sys.stdout.flush()
    try:
        srv.serve_forever()
    except KeyboardInterrupt:
        srv.shutdown()