import argparse
import copy
import datetime
import hmac
import json
import os
import sys
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[2]
DEFAULT_FIXTURE = REPO_ROOT / "docs" / "agent-api" / "summary.v1.json"
LOCK = threading.Lock()
MESH_PRICE = 3.6
ROUTES = {
"/v1/summary": {"GET"},
"/v1/refresh": {"POST"},
"/v1/sharing": {"PUT"},
"/v1/workers": {"PUT"},
"/v1/price": {"PUT"},
}
PRICE_FIELDS = {"scope", "price_per_hour"}
def now_iso():
return datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
def write_agent_json(home, port, token):
agent_dir = Path(home).expanduser() / "agent"
agent_dir.mkdir(parents=True, exist_ok=True)
os.chmod(str(agent_dir), 0o700)
path = agent_dir / "agent.json"
path.write_text(json.dumps({"v": 1, "port": port, "token": token, "pid": os.getpid(), "version": "stub"}))
os.chmod(str(path), 0o600)
return path
class State:
def __init__(self, summary, fail_price):
self.s = summary
self.fail_price = fail_price
def this_mac_device(self):
for device in self.s.get("devices") or []:
if device.get("this_mac"):
return device
return None
def apply(self, path, payload):
mac = self.s["this_mac"]
workers = mac["workers"]
if path == "/v1/sharing":
if not isinstance(payload.get("on"), bool):
return (400, "invalid_input", "on must be a boolean")
mac["sharing"] = payload["on"]
mac["broker"] = "running" if payload["on"] else "stopped"
workers["running"] = workers["desired"] if payload["on"] else 0
workers["busy"] = min(workers["busy"], workers["running"])
return None
if path == "/v1/workers":
count = payload.get("count")
if not isinstance(count, int) or isinstance(count, bool) or not 0 <= count <= workers["max"]:
return (400, "invalid_input", "count must be an integer in 0..%d" % workers["max"])
workers["desired"] = count
if mac["sharing"]:
workers["running"] = count
workers["busy"] = min(workers["busy"], count)
device = self.this_mac_device()
if device is not None:
device["workers"] = count
return None
if path == "/v1/price":
if self.fail_price:
return (502, "hub_error", "stub: the hub rejected the price")
if not set(payload.keys()) <= PRICE_FIELDS:
return (400, "invalid_input", "unknown field in price body")
scope = payload.get("scope")
if scope not in ("device", "default"):
return (400, "invalid_input", "scope must be device or default")
if "price_per_hour" not in payload:
return (400, "invalid_input", "price_per_hour is required (null falls back)")
value = payload["price_per_hour"]
bounds = self.s["prices"]["bounds"]
if value is not None and (isinstance(value, bool) or not isinstance(value, (int, float))
or not bounds["min"] <= value <= bounds["max"]):
return (400, "invalid_input", "price_per_hour must be within %s..%s" % (bounds["min"], bounds["max"]))
prices = self.s["prices"]
price = prices["this_mac"]
if scope == "default":
prices["default_per_hour"] = value
else:
price["per_hour"] = value
price["state"] = "set" if value is not None else "inherited"
fallback = prices["default_per_hour"] if prices["default_per_hour"] is not None else MESH_PRICE
price["effective_per_hour"] = price["per_hour"] if price["state"] == "set" else fallback
device = self.this_mac_device()
if device is not None:
device["effective_price_per_hour"] = {"min": price["effective_per_hour"], "max": price["effective_per_hour"]}
return None
return (404, "not_found", path)
def make_handler(state, port, token):
allowed_hosts = {"127.0.0.1:%d" % port, "localhost:%d" % port}
class Handler(BaseHTTPRequestHandler):
server_version = "zc-agent-stub"
def log_message(self, fmt, *args):
sys.stderr.write("stub-agent: %s %s\n" % (self.command, self.path))
def reply(self, status, body):
data = json.dumps(body).encode()
self.send_response(status)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(data)))
self.end_headers()
self.wfile.write(data)
def fail(self, status, code, message):
self.reply(status, {"error": {"code": code, "message": message}})
def guard(self):
if self.headers.get("Origin") is not None:
self.fail(403, "forbidden", "requests with an Origin header are refused")
return False
if self.headers.get("Host") not in allowed_hosts:
self.fail(403, "forbidden", "Host must be 127.0.0.1:%d or localhost:%d" % (port, port))
return False
header = self.headers.get("Authorization") or ""
if not hmac.compare_digest(header, "Bearer " + token):
self.fail(401, "unauthorized", "missing or invalid token")
return False
return True
def summary(self):
with LOCK:
state.s["as_of"] = now_iso()
return copy.deepcopy(state.s)
def handle_put(self, path):
length = int(self.headers.get("Content-Length") or 0)
try:
payload = json.loads(self.rfile.read(length) or b"{}")
except ValueError:
payload = None
if not isinstance(payload, dict):
self.fail(400, "invalid_input", "body must be a JSON object")
return
with LOCK:
error = state.apply(path, payload)
if error is not None:
self.fail(*error)
return
self.reply(200, self.summary())
def route(self):
if not self.guard():
return
path = self.path
methods = ROUTES.get(path)
if methods is None:
self.fail(404, "not_found", path)
return
if self.command not in methods:
self.fail(405, "method_not_allowed", "%s not allowed on %s" % (self.command, path))
return
if self.command == "PUT":
self.handle_put(path)
else:
self.reply(200, self.summary())
do_GET = do_POST = do_PUT = do_DELETE = do_PATCH = do_HEAD = do_OPTIONS = route
return Handler
def main():
parser = argparse.ArgumentParser(description="Stub of the zc agent local API v1")
parser.add_argument("--port", type=int, default=4720)
parser.add_argument("--home", default="~/.zakuro-stub")
parser.add_argument("--token", default="stub-token")
parser.add_argument("--fixture", default=str(DEFAULT_FIXTURE))
parser.add_argument("--fail-price", action="store_true", help="answer PUT /v1/price with 502 hub_error")
parser.add_argument("--summary-version", type=int, default=1, help="serve this v (2 exercises 'Update zc')")
parser.add_argument("--problem", action="append", default=None, help="replace problems with these codes")
args = parser.parse_args()
summary = json.loads(Path(args.fixture).read_text())
summary["v"] = args.summary_version
if args.problem is not None:
summary["problems"] = [
{"code": code, "message": "stub: %s" % code, "hint": "stub hint for %s" % code, "detail": None}
for code in args.problem
]
path = write_agent_json(args.home, args.port, args.token)
handler = make_handler(State(summary, args.fail_price), args.port, args.token)
server = ThreadingHTTPServer(("127.0.0.1", args.port), handler)
print("stub agent on http://127.0.0.1:%d (token %r); agent.json at %s" % (args.port, args.token, path))
print("point the app at it: defaults write ai.zakuro.Zakuro ZakuroHome %s" % Path(args.home).expanduser())
try:
server.serve_forever()
except KeyboardInterrupt:
pass
finally:
server.server_close()
if __name__ == "__main__":
main()