277 lines
14 KiB
Python
277 lines
14 KiB
Python
"""Фейковый хаб для тестов агента: маршруты /agent/* в памяти."""
|
||
import base64
|
||
import hashlib
|
||
import json
|
||
import re
|
||
import socket
|
||
import threading
|
||
import time
|
||
from datetime import datetime
|
||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||
|
||
|
||
MAX_PAGES = 50
|
||
MAX_FILE = 10 * 1024 * 1024
|
||
|
||
|
||
def check_pages(command, command_id, pages):
|
||
"""Форма пачки как в stage.accept/_check хаба: (код, тело ошибки) или None. 422 — {"error", "detail"}."""
|
||
def bad(detail, status=422, error="bad_pages"):
|
||
return status, {"error": error, "detail": detail}
|
||
|
||
if isinstance(command_id, bool) or not isinstance(command_id, int):
|
||
return bad("command_id — целое", error="bad_request")
|
||
if not isinstance(pages, list) or not pages or len(pages) > MAX_PAGES:
|
||
return bad("pages — список от 1 до %d страниц" % MAX_PAGES)
|
||
if command is None:
|
||
return bad("command %s" % command_id, 404, "not_found")
|
||
if command["type"] != "wiki_sync" or command["status"] != "claimed":
|
||
return bad("команда %s не выдана агенту как wiki_sync" % command_id, 409, "not_claimed")
|
||
limits = {"code": 200, "title": 1000, "url": 4000, "parent_code": 200, "eva_id": 200}
|
||
for i, page in enumerate(pages, 1):
|
||
if not isinstance(page, dict):
|
||
return bad("страница %d: не объект" % i)
|
||
code = page.get("code")
|
||
if not isinstance(code, str) or not code.strip() or len(code) > limits["code"]:
|
||
return bad("страница %d: code" % i)
|
||
for key, required in (("title", True), ("parent_code", False), ("url", False)):
|
||
v = page.get(key)
|
||
if v is None and not required:
|
||
continue
|
||
if not isinstance(v, str) or (required and not v.strip()) or len(v) > limits[key]:
|
||
return bad("%s: %s" % (code, key))
|
||
if not isinstance(page.get("html"), str):
|
||
return bad("%s: html" % code)
|
||
o = page.get("orderno")
|
||
if o is not None and (not isinstance(o, int) or isinstance(o, bool)):
|
||
return bad("%s: orderno" % code)
|
||
e = page.get("eva_id")
|
||
if e is not None and (not isinstance(e, str) or len(e) > limits["eva_id"]):
|
||
return bad("%s: eva_id" % code)
|
||
att, links = page.get("attachments") or [], page.get("links") or []
|
||
if not isinstance(att, list) or not all(
|
||
isinstance(a, dict) and isinstance(a.get("ref"), str) and isinstance(a.get("url"), str)
|
||
and isinstance(a.get("name") or "", str) for a in att):
|
||
return bad("%s: attachments" % code)
|
||
if not isinstance(links, list) or not all(
|
||
isinstance(l, dict) and isinstance(l.get("ref"), str) and isinstance(l.get("code"), str)
|
||
for l in links):
|
||
return bad("%s: links" % code)
|
||
return None
|
||
|
||
|
||
def check_versions(command, command_id, versions):
|
||
"""Форма пачки версий как в приёме версий хаба (change wiki-history): (код, тело ошибки) или None."""
|
||
def bad(detail, status=422, error="bad_versions"):
|
||
return status, {"error": error, "detail": detail}
|
||
|
||
if isinstance(command_id, bool) or not isinstance(command_id, int):
|
||
return bad("command_id — целое", error="bad_request")
|
||
if not isinstance(versions, list) or not versions or len(versions) > MAX_PAGES:
|
||
return bad("versions — список от 1 до %d версий" % MAX_PAGES)
|
||
if command is None:
|
||
return bad("command %s" % command_id, 404, "not_found")
|
||
if command["type"] != "wiki_sync" or command["status"] != "claimed":
|
||
return bad("команда %s не выдана агенту как wiki_sync" % command_id, 409, "not_claimed")
|
||
for i, v in enumerate(versions, 1):
|
||
code = v.get("code") if isinstance(v, dict) else None
|
||
if not isinstance(code, str) or not code.strip():
|
||
return bad("версия %d: code" % i)
|
||
page = v.get("page_code")
|
||
if not isinstance(page, str) or not page.strip():
|
||
return bad("%s: page_code" % code)
|
||
try:
|
||
ok = isinstance(v.get("created_at"), str) and datetime.fromisoformat(v["created_at"]).tzinfo is not None
|
||
except ValueError:
|
||
ok = False
|
||
if not ok:
|
||
return bad("%s: created_at" % code)
|
||
a = v.get("author")
|
||
if not isinstance(a, dict) or not all(a.get(k) is None or isinstance(a.get(k), str) for k in ("login", "name")):
|
||
return bad("%s: author" % code)
|
||
if not isinstance(v.get("html"), str):
|
||
return bad("%s: html" % code)
|
||
att, links = v.get("attachments"), v.get("links")
|
||
if not isinstance(att, list) or not all(
|
||
isinstance(x, dict) and isinstance(x.get("ref"), str) and isinstance(x.get("url"), str)
|
||
and isinstance(x.get("name") or "", str) for x in att):
|
||
return bad("%s: attachments" % code)
|
||
if not isinstance(links, list) or not all(
|
||
isinstance(x, dict) and isinstance(x.get("ref"), str) and isinstance(x.get("code"), str)
|
||
for x in links):
|
||
return bad("%s: links" % code)
|
||
return None
|
||
|
||
|
||
class FakeHub:
|
||
token = "agent-token"
|
||
|
||
def __init__(self, settings):
|
||
self.settings = dict(settings)
|
||
self.queue = []
|
||
self.commands = {}
|
||
self.results = {}
|
||
self.claims = 0
|
||
self.claim_bodies = []
|
||
self.fail_results = 0
|
||
self.reject_results = {}
|
||
self.in_flight = None
|
||
self.release = None # {"version", "sha256", "content"} — опубликованная версия агента
|
||
self.release_requests = 0
|
||
self.files = {} # sha256 → содержимое (хранилище файлов хаба)
|
||
self.file_requests = 0
|
||
self.pages = [] # тела POST /agent/wiki/pages
|
||
self.reject_pages = None # (код, тело) — ответ на /agent/wiki/pages ошибкой
|
||
self.versions = [] # тела принятых POST /agent/wiki/versions
|
||
self.reject_versions = None # (код, тело) — ответ на /agent/wiki/versions ошибкой
|
||
self.fail_versions = 0 # сколько ближайших /agent/wiki/versions ответить 503
|
||
self.wiki_posts = [] # пути POST /agent/wiki/pages|versions по порядку (и отклонённых)
|
||
self.claim_wait = 0 # с: при пустой очереди держать claim, как хаб
|
||
self.uploaded = [] # содержимое, присланное агентом POST /agent/files
|
||
self.reject_files = None # (код, тело) — ответ на POST /agent/files ошибкой; "drop" — оборвать без ответа
|
||
self.lock = threading.Lock()
|
||
hub = self
|
||
|
||
class Handler(BaseHTTPRequestHandler):
|
||
def log_message(self, *a):
|
||
pass
|
||
|
||
def do_GET(self):
|
||
hub.handle(self, "GET")
|
||
|
||
def do_POST(self):
|
||
hub.handle(self, "POST")
|
||
|
||
self.server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
|
||
threading.Thread(target=self.server.serve_forever, daemon=True).start()
|
||
|
||
@property
|
||
def base_url(self):
|
||
return "http://127.0.0.1:%d" % self.server.server_port
|
||
|
||
def stop(self):
|
||
self.server.shutdown()
|
||
self.server.server_close()
|
||
|
||
def enqueue(self, ctype, params, pause=0):
|
||
cid = len(self.commands) + 1
|
||
self.commands[cid] = {"id": cid, "type": ctype, "params": params, "pause_seconds": pause, "status": "queued"}
|
||
self.queue.append(cid)
|
||
return cid
|
||
|
||
def put_file(self, data):
|
||
sha = hashlib.sha256(data).hexdigest()
|
||
self.files[sha] = data
|
||
return sha
|
||
|
||
def _issued(self, c):
|
||
return dict(c["params"])
|
||
|
||
def _release_ref(self):
|
||
r = self.release
|
||
return {"version": r["version"], "sha256": r["sha256"]} if r else None
|
||
|
||
def handle(self, h, method):
|
||
length = int(h.headers.get("Content-Length") or 0)
|
||
body = json.loads(h.rfile.read(length).decode("utf-8")) if length else {}
|
||
if h.headers.get("Authorization") != "Bearer " + self.token:
|
||
return self._send(h, 401, {"error": "unauthorized"})
|
||
path = h.path.split("?")[0]
|
||
if path == "/agent/files" and method == "POST" and self.reject_files == "drop":
|
||
h.close_connection = True
|
||
h.connection.shutdown(socket.SHUT_RDWR)
|
||
return
|
||
if path == "/agent/claim" and self.claim_wait:
|
||
end = time.monotonic() + self.claim_wait
|
||
while time.monotonic() < end: # ждать без блокировки: другие запросы (enqueue) проходят
|
||
with self.lock:
|
||
if self.queue:
|
||
break
|
||
time.sleep(0.05)
|
||
with self.lock:
|
||
code, payload = self.route(method, path, body)
|
||
self._send(h, code, payload)
|
||
|
||
def _send(self, h, code, payload):
|
||
data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
|
||
h.send_response(code)
|
||
h.send_header("Content-Type", "application/json")
|
||
h.send_header("Content-Length", str(len(data)))
|
||
h.end_headers()
|
||
h.wfile.write(data)
|
||
|
||
def route(self, method, path, body):
|
||
if path == "/agent/claim":
|
||
self.claim_bodies.append(body or {})
|
||
self.claims += 1
|
||
if self.in_flight and self.commands[self.in_flight]["status"] == "claimed":
|
||
return 200, {"command": None, "settings_version": self.settings["settings_version"],
|
||
"agent_release": self._release_ref()}
|
||
if not self.queue:
|
||
return 200, {"command": None, "settings_version": self.settings["settings_version"],
|
||
"agent_release": self._release_ref()}
|
||
cid = self.queue.pop(0)
|
||
self.commands[cid]["status"] = "claimed"
|
||
self.in_flight = cid
|
||
c = self.commands[cid]
|
||
return 200, {"command": {"id": cid, "type": c["type"], "params": self._issued(c), "pause_seconds": c["pause_seconds"]},
|
||
"settings_version": self.settings["settings_version"], "agent_release": self._release_ref()}
|
||
if path == "/agent/release":
|
||
self.release_requests += 1
|
||
return (200, self.release) if self.release else (404, {"error": "no_release"})
|
||
if path == "/agent/settings":
|
||
return 200, self.settings
|
||
if path == "/agent/wiki/versions" and method == "POST":
|
||
self.wiki_posts.append(path)
|
||
if self.fail_versions > 0:
|
||
self.fail_versions -= 1
|
||
return 503, {"error": "busy"}
|
||
if self.reject_versions:
|
||
return self.reject_versions
|
||
cid = body.get("command_id")
|
||
bad = check_versions(self.commands.get(cid) if isinstance(cid, int) else None, cid, body.get("versions"))
|
||
if bad:
|
||
return bad
|
||
self.versions.append(body)
|
||
return 200, {"accepted": len(body.get("versions") or [])}
|
||
if path == "/agent/wiki/pages" and method == "POST":
|
||
self.wiki_posts.append(path)
|
||
if self.reject_pages:
|
||
return self.reject_pages
|
||
cid = body.get("command_id")
|
||
bad = check_pages(self.commands.get(cid) if isinstance(cid, int) else None, cid, body.get("pages"))
|
||
if bad:
|
||
return bad
|
||
self.pages.append(body)
|
||
return 200, {"accepted": len(body.get("pages") or [])}
|
||
if path == "/agent/files" and method == "POST":
|
||
if self.reject_files:
|
||
return self.reject_files
|
||
data = base64.b64decode(body.get("content") or "", validate=True)
|
||
if len(data) > MAX_FILE: # как files.put хаба; пустой файл от агента принимается (allow_empty)
|
||
return 413, {"error": "too_large", "detail": "файл больше %d байт (%d)" % (MAX_FILE, len(data))}
|
||
self.uploaded.append(data)
|
||
return 200, {"file": self.put_file(data), "size": len(data)}
|
||
m = re.match(r"^/agent/commands/(\d+)$", path)
|
||
if m:
|
||
c = self.commands.get(int(m.group(1)))
|
||
return (200, dict(c, params=self._issued(c))) if c else (404, {"error": "not_found"})
|
||
m = re.match(r"^/agent/commands/(\d+)/result$", path)
|
||
if m:
|
||
cid = int(m.group(1))
|
||
if self.fail_results > 0:
|
||
self.fail_results -= 1
|
||
return 503, {"error": "busy"}
|
||
if cid in self.reject_results:
|
||
return self.reject_results[cid], {"error": "rejected"}
|
||
if cid not in self.commands:
|
||
return 404, {"error": "not_found"}
|
||
if self.commands[cid]["status"] in ("done", "failed", "cancelled"):
|
||
return 200, self.commands[cid] # терминальный статус не перезаписывается (как живой хаб)
|
||
self.results[cid] = body
|
||
self.commands[cid]["status"] = body.get("status")
|
||
if self.in_flight == cid:
|
||
self.in_flight = None
|
||
return 200, self.commands[cid]
|
||
return 404, {"error": "not_found", "detail": path}
|