harden probe agent v1.0.3

安全鉴权、默认只读、显式提交开关、状态接口与测试
2026-08-26 17:16:02 +08:00
parent 5db73110a5
commit f18eeaf334
4 changed files with 368 additions and 91 deletions
--- a/main.py
+++ b/main.py
@@ -1,15 +1,10 @@
-"""ModelHub XC 适配智能体 · 只读探针 + 容器内提交测试 v1.0.2
+"""ModelHub XC 适配智能体安全探针 v1.0.3。

-关键实验:在容器里用注入的 EXTERNAL_SERVICE_TOKEN(=tmp) + STRATEGY_ID 真调一次 task/add,
-验证平台是否靠「请求来自容器 pod」放行智能体提交(浏览器/脚本提交会被 60014 挡)。
-
-诊断输出:
-  - dump 所有 env 的名字+长度(不打值)。
-  - EXTERNAL_SERVICE_TOKEN 值(短则原样、长则打码)。
-  - 容器内 build-config + task/add(真实模型) 的返回 code/msg。
-
-约束:token 值不打(短占位除外);只提交这一次测试(真实模型可能进队列);SIGTERM 优雅停机。
+默认只做一次只读鉴权检查。只有显式设置 ALLOW_SUBMIT=1，且鉴权检查成功后，
+才会调用 build-config 和 task/add；每个进程生命周期最多提交一次。
 """
+
+import copy
 import json
 import os
 import signal
@@ -18,100 +13,234 @@ import time
 import urllib.error
 import urllib.request
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+from typing import Callable, Mapping, Optional, Tuple

-MAIN = os.getenv("MAIN_HOST", "https://modelhub.org.cn")
-TOKEN = os.getenv("EXTERNAL_SERVICE_TOKEN", "")
-STRATEGY_ID = os.getenv("STRATEGY_ID", "")
+
+VERSION = "1.0.3"
+MAIN = os.getenv("MAIN_HOST", "https://modelhub.org.cn").rstrip("/")
+STRATEGY_ID = os.getenv("STRATEGY_ID", "").strip()
 PORT = int(os.getenv("PORT", "8080"))
+ALLOW_SUBMIT = os.getenv("ALLOW_SUBMIT", "").strip().lower() in {"1", "true", "yes", "on"}

-TEST_MODEL = os.getenv("TEST_MODEL", "https://www.modelscope.cn/models/mradermacher/TrialSpace-1225-GGUF")
-GPU, FW, TT = os.getenv("GPU", "MetaX_c-500"), os.getenv("FW", "vllm"), os.getenv("TT", "text-generation")
-
-HARDCODED_CFG = (
-    "framework: vllm\napi: completion\nlang: en\nmax_model_len: 4096\nmax_tokens: 1024\n"
-    "temperature: 0.7\nrepetition_penalty: 1.1\ntop_p: 0.9\nsut_config:\n  gpu_num: 1\n  values:\n"
-    "    command:\n    - /opt/conda/bin/vllm\n    - serve\n    - /model\n    - --port\n    - '20644'\n"
-    "    - --served-model-name\n    - llm\n    - --max-model-len\n    - '4096'\n"
-    "    - --gpu-memory-utilization\n    - '0.9'\n    - -tp\n    - '1'\n    - --enforce-eager\n"
-    "    - --trust-remote-code\nref_config:\n  gpu_num: 1\n  values:\n    command:\n    - vllm\n"
-    "    - serve\n    - /model\n    - --port\n    - '80'\n    - --served-model-name\n    - llm\n"
-    "    - --max-model-len\n    - '4096'\n    - -tp\n    - '1'\n    - --enforce-eager\n    - --trust-remote-code\n"
+TEST_MODEL = os.getenv(
+    "TEST_MODEL",
+    "https://www.modelscope.cn/models/mradermacher/TrialSpace-1225-GGUF",
 )
+GPU = os.getenv("GPU", "MetaX_c-500")
+FW = os.getenv("FW", "vllm")
+TT = os.getenv("TT", "text-generation")

+_TOKEN_PLACEHOLDERS = {"tmp", "placeholder", "changeme", "change-me", "example", "test"}
 shutdown = threading.Event()
-PROBE = {"done": False, "results": []}
+STATE_LOCK = threading.Lock()
+
+
+def _initial_state() -> dict:
+    return {
+        "version": VERSION,
+        "phase": "starting",
+        "done": False,
+        "auth_ready": False,
+        "submit_enabled": ALLOW_SUBMIT,
+        "token_source": "none",
+        "results": {},
+    }
+
+
+STATE = _initial_state()


 def log(msg: str) -> None:
    print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {msg}", flush=True)


-def _mask(v: str) -> str:
-    return v if len(v) <= 8 else (v[:2] + "…" + v[-2:])
+def _set_state(**changes) -> None:
+    with STATE_LOCK:
+        STATE.update(changes)


-def _req(method: str, path: str, headers: dict, body=None, timeout: int = 30) -> dict:
-    data = body.encode() if isinstance(body, str) else body
+def _snapshot() -> dict:
+    with STATE_LOCK:
+        return copy.deepcopy(STATE)
+
+
+def _reset_state() -> None:
+    """Reset process state; kept separate to make the probe deterministically testable."""
+    with STATE_LOCK:
+        STATE.clear()
+        STATE.update(_initial_state())
+
+
+def _valid_token(value: str) -> bool:
+    value = value.strip()
+    return bool(value) and value.lower() not in _TOKEN_PLACEHOLDERS
+
+
+def _resolve_token(
+    env: Optional[Mapping[str, str]] = None,
+    read_text: Optional[Callable[[str], str]] = None,
+) -> Tuple[str, str, Optional[str]]:
+    """Resolve a credential without exposing its value or length.
+
+    Precedence is XC_TOKEN_FILE, XC_TOKEN, then the legacy
+    EXTERNAL_SERVICE_TOKEN. An explicitly configured but unreadable token file
+    is treated as an error instead of silently falling back.
+    """
+    values = os.environ if env is None else env
+    file_reader = read_text or (lambda path: Path(path).read_text(encoding="utf-8"))
+    token_file = values.get("XC_TOKEN_FILE", "").strip()
+    if token_file:
+        try:
+            value = file_reader(token_file).strip()
+        except (OSError, UnicodeError):
+            return "", "XC_TOKEN_FILE", "token_file_unreadable"
+        if _valid_token(value):
+            return value, "XC_TOKEN_FILE", None
+        return "", "XC_TOKEN_FILE", "token_missing_or_placeholder"
+
+    for name in ("XC_TOKEN", "EXTERNAL_SERVICE_TOKEN"):
+        value = values.get(name, "").strip()
+        if _valid_token(value):
+            return value, name, None
+        if value:
+            return "", name, "token_missing_or_placeholder"
+    return "", "none", "token_missing_or_placeholder"
+
+
+def _req(method: str, path: str, token: str, body=None, timeout: int = 30) -> dict:
+    headers = {"Accept": "application/json", "Xc-Token": token}
+    data = body.encode("utf-8") if isinstance(body, str) else body
+    if data is not None:
+        headers["Content-Type"] = "application/json"
    req = urllib.request.Request(MAIN + path, data=data, method=method, headers=headers)
    try:
-        with urllib.request.urlopen(req, timeout=timeout) as r:
-            txt = r.read().decode("utf-8", "ignore")
-            status = r.status
-    except urllib.error.HTTPError as e:
-        txt = e.read().decode("utf-8", "ignore")
-        status = e.code
-    except Exception as e:  # noqa: BLE001
-        return {"http": -1, "err": str(e)[:100]}
+        with urllib.request.urlopen(req, timeout=timeout) as response:
+            text = response.read().decode("utf-8", "ignore")
+            status = response.status
+    except urllib.error.HTTPError as error:
+        text = error.read().decode("utf-8", "ignore")
+        status = error.code
+    except Exception as error:  # noqa: BLE001 - network diagnostics are returned safely
+        return {"http": -1, "error": type(error).__name__}
+
    try:
-        j = json.loads(txt)
-        return {"http": status, "code": j.get("code"), "msg": (j.get("message") or "")[:100],
-                "data": j.get("data")}
-    except Exception:  # noqa: BLE001
-        return {"http": status, "raw": txt[:100]}
+        payload = json.loads(text)
+    except (TypeError, ValueError):
+        return {"http": status, "code": None, "message": "non_json_response"}
+    return {
+        "http": status,
+        "code": payload.get("code"),
+        "message": str(payload.get("message") or "")[:100],
+        "data": payload.get("data"),
+    }
+
+
+def _api_ok(result: dict) -> bool:
+    try:
+        status = int(result.get("http", -1))
+    except (TypeError, ValueError):
+        return False
+    return 200 <= status < 300 and result.get("code") in {0, "0"}
+
+
+def _result_summary(result: dict) -> dict:
+    summary = {"http": result.get("http"), "code": result.get("code")}
+    if result.get("message"):
+        summary["message"] = result["message"]
+    if result.get("error"):
+        summary["error"] = result["error"]
+    return summary
+
+
+def _task_id(data) -> Optional[str]:
+    if not isinstance(data, dict):
+        return None
+    for key in ("id", "taskId", "task_id"):
+        if data.get(key) is not None:
+            return str(data[key])
+    return None


 def probe() -> None:
-    env_lens = {k: len(str(os.environ.get(k, ""))) for k in sorted(os.environ)}
-    log("env keys+lens: " + json.dumps(env_lens, ensure_ascii=False))
-    log(f"strategy_id={STRATEGY_ID or 'MISSING'} | EXTERNAL_SERVICE_TOKEN value={_mask(TOKEN)!r} len={len(TOKEN)}")
+    """Run one guarded probe; no retry loop and at most one task/add call."""
+    token, source, token_error = _resolve_token()
+    _set_state(token_source=source, submit_enabled=ALLOW_SUBMIT)

-    # 只读:token 有效性
-    page = _req("GET", "/api/adapt/task/page?current=1&pageSize=1", {"Xc-Token": TOKEN})
-    log("EST-token task/page -> " + json.dumps({k: page.get(k) for k in ("http", "code")}, ensure_ascii=False))
+    if not STRATEGY_ID:
+        _set_state(phase="blocked_missing_strategy_id", done=True)
+        log("probe blocked: STRATEGY_ID is missing")
+        return
+    if token_error:
+        _set_state(phase="blocked_missing_auth", done=True, results={"auth": token_error})
+        log(f"probe blocked: usable credential is unavailable (source={source})")
+        return

-    # 关键:容器内 build-config + task/add(真实模型)
-    bc = _req("POST", f"/api/adapt/task/build-config?gpuType={GPU}&framework={FW}&taskType={TT}", {"Xc-Token": TOKEN})
-    log("build-config(EST-token) -> " + json.dumps({k: bc.get(k) for k in ("http", "code", "msg")}, ensure_ascii=False))
-    cfg = bc.get("data") if isinstance(bc.get("data"), str) and bc.get("data") else HARDCODED_CFG
-    body = json.dumps({"modelAddress": TEST_MODEL, "taskType": TT, "targetGpu": GPU,
-                       "framework": FW, "strategyId": STRATEGY_ID, "configParams": cfg})
-    ta = _req("POST", "/api/adapt/task/add", {"Xc-Token": TOKEN, "Content-Type": "application/json"}, body)
-    log("SUBMIT-TEST task/add(EST-token+strategyId, real model) -> " + json.dumps(ta, ensure_ascii=False))
-    PROBE["results"] = [{"task_page": page.get("code")}, {"build_config": bc.get("code")},
-                        {"task_add": ta.get("code"), "msg": ta.get("msg")}]
-    PROBE["done"] = True
-    log("probe done | task/add code=0 => 容器提交通; 60014 => 容器来源也被判脚本; 401 => token 无效")
+    _set_state(phase="checking_auth")
+    page = _req("GET", "/api/adapt/task/page?current=1&pageSize=1", token)
+    results = {"task_page": _result_summary(page)}
+    if not _api_ok(page):
+        _set_state(phase="auth_failed", done=True, results=results)
+        log("read-only auth check failed: " + json.dumps(results["task_page"], ensure_ascii=False))
+        return
+
+    _set_state(auth_ready=True, results=results)
+    log(f"read-only auth check passed (source={source})")
+    if not ALLOW_SUBMIT:
+        _set_state(phase="read_only_complete", done=True)
+        log("probe complete in read-only mode; task/add was not called")
+        return
+
+    _set_state(phase="preparing_submission")
+    build_path = f"/api/adapt/task/build-config?gpuType={GPU}&framework={FW}&taskType={TT}"
+    build_config = _req("POST", build_path, token)
+    results["build_config"] = _result_summary(build_config)
+    config = build_config.get("data")
+    if not _api_ok(build_config) or not isinstance(config, str) or not config:
+        _set_state(phase="build_config_failed", done=True, results=results)
+        log("build-config failed; task/add was not called")
+        return
+
+    body = json.dumps(
+        {
+            "modelAddress": TEST_MODEL,
+            "taskType": TT,
+            "targetGpu": GPU,
+            "framework": FW,
+            "strategyId": STRATEGY_ID,
+            "configParams": config,
+        }
+    )
+    task_add = _req("POST", "/api/adapt/task/add", token, body)
+    results["task_add"] = _result_summary(task_add)
+    task_id = _task_id(task_add.get("data"))
+    if task_id:
+        results["task_add"]["task_id"] = task_id
+
+    phase = "submit_complete" if _api_ok(task_add) else "submit_failed"
+    _set_state(phase=phase, done=True, results=results)
+    log("single task/add attempt finished: " + json.dumps(results["task_add"], ensure_ascii=False))


 class Handler(BaseHTTPRequestHandler):
    def _json(self, obj: dict, status: int = 200) -> None:
-        payload = json.dumps(obj, ensure_ascii=False).encode()
+        payload = json.dumps(obj, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
-        self.send_header("Content-Type", "application/json")
+        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(payload)))
        self.end_headers()
        self.wfile.write(payload)

    def do_GET(self) -> None:
        if self.path == "/health":
-            self._json({"status": "ok"})
+            self._json({"status": "ok", "version": VERSION})
            return
-        if self.path == "/":
-            self._json({"name": "huni-probe-agent", "mode": "submit-test-v1.0.2", "probe": PROBE})
+        if self.path in {"/", "/status"}:
+            self._json({"name": "huni-probe-agent", "probe": _snapshot()})
            return
        self._json({"error": "not found"}, 404)

-    def log_message(self, *_a) -> None:
+    def log_message(self, *_args) -> None:
        pass


@@ -127,7 +256,7 @@ def main() -> None:

    server = ThreadingHTTPServer(("0.0.0.0", PORT), Handler)
    server.timeout = 1
-    log(f"probe agent v1.0.2 listening on 0.0.0.0:{PORT}")
+    log(f"probe agent v{VERSION} listening on 0.0.0.0:{PORT} (submit_enabled={ALLOW_SUBMIT})")
    while not shutdown.is_set():
        server.handle_request()
    server.server_close()