268 lines
8.9 KiB
Python
268 lines
8.9 KiB
Python
"""ModelHub XC 适配智能体安全探针 v1.0.3。
|
||
|
||
默认只做一次只读鉴权检查。只有显式设置 ALLOW_SUBMIT=1，且鉴权检查成功后，
|
||
才会调用 build-config 和 task/add；每个进程生命周期最多提交一次。
|
||
"""
|
||
|
||
import copy
|
||
import json
|
||
import os
|
||
import signal
|
||
import threading
|
||
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
|
||
|
||
|
||
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 = 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()
|
||
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 _set_state(**changes) -> None:
|
||
with STATE_LOCK:
|
||
STATE.update(changes)
|
||
|
||
|
||
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 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:
|
||
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:
|
||
"""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)
|
||
|
||
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
|
||
|
||
_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("utf-8")
|
||
self.send_response(status)
|
||
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", "version": VERSION})
|
||
return
|
||
if self.path in {"/", "/status"}:
|
||
self._json({"name": "huni-probe-agent", "probe": _snapshot()})
|
||
return
|
||
self._json({"error": "not found"}, 404)
|
||
|
||
def log_message(self, *_args) -> None:
|
||
pass
|
||
|
||
|
||
def _handle_signal(signum: int, _frame) -> None:
|
||
shutdown.set()
|
||
log(f"received signal {signum}, shutting down")
|
||
|
||
|
||
def main() -> None:
|
||
signal.signal(signal.SIGTERM, _handle_signal)
|
||
signal.signal(signal.SIGINT, _handle_signal)
|
||
threading.Thread(target=probe, daemon=True).start()
|
||
|
||
server = ThreadingHTTPServer(("0.0.0.0", PORT), Handler)
|
||
server.timeout = 1
|
||
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()
|
||
log("stopped")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|