cc2-spoolman-bridge/tools/mock_env.py

"""Local stand-in for the printer (SSH), CC2 Control and Spoolman, for manual end-to-end tests.

    python tools/mock_env.py --host mock --printers 2

Per printer i: SSH 2222+i (keys added via POST /authorize), CC2 Control 8081+i. Spoolman 7912, control 8099.
Control endpoints (port 8099, optional &printer=i): POST /insert?slot=&spool=, /remove?slot=,
/consume?slot=&grams=, /job?active=0|1, /touch?slot= (touchscreen edit: weight falls back to 1000).
"""

from __future__ import annotations

import argparse
import itertools
import json
import socket
import sys
import threading
from pathlib import Path

import paramiko
import uvicorn
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from tests.fakes import channel, record  # noqa: E402

def printer_state():
    st = {
        "record": record(),
        "enabled": True,
        "slots": [{"slot": i, "spool": "", "last": "", "question": "", "printer": {"status": 0}} for i in range(5)],
        "spools": [],
        "job": {"active": False, "file": "", "slots": []},
        "ids": (f"m{i:04d}" for i in itertools.count(1)),
    }
    st["slots"][4]["printer"] = None
    return st


PRINTERS: list[dict] = []
STATE = {"spoolman": {}, "fields": {}, "requests": []}
LOCK = threading.Lock()


def seed_spoolman():
    for sid, vendor, mat, color, rem in [(1, "Sunlu", "PLA+", "000000", 640.0), (5, "Sunlu", "PLA+", "1E40AF", 812.5),
                                         (8, "Elegoo", "PLA", "FF7F00", 990.0), (12, "Elegoo", "PETG", "F5F5F5", 455.0),
                                         (14, "Polymaker", "PLA", "C0A060", 300.0)]:
        STATE["spoolman"][sid] = {
            "id": sid, "lot_nr": f"{vendor.upper()}{sid:03d}", "remaining_weight": rem, "used_weight": 1000 - rem,
            "initial_weight": 1000.0, "spool_weight": 160.0, "location": "Lager", "archived": False, "price": 19.9,
            "extra": {},
            "filament": {"id": sid, "name": f"{mat} {'Silk Gold' if sid == 14 else color}", "material": mat,
                         "color_hex": color, "diameter": 1.75, "density": 1.24, "weight": 1000.0,
                         "settings_extruder_temp": 215, "vendor": {"id": 1, "name": vendor}},
        }


# ---------- CC2 Control ----------

def make_cc2(host_header: str, P: dict) -> FastAPI:
    app = FastAPI()

    @app.middleware("http")
    async def guard(request: Request, call_next):
        raw = [k.decode().lower() for k, _ in request.scope["headers"]]
        if len(raw) != len(set(raw)):
            return JSONResponse({"error": "duplicate header"}, 403)
        if request.headers.get("host") != host_header or "origin" in request.headers \
                or "transfer-encoding" in request.headers:
            return JSONResponse({"error": "guard"}, 403)
        if request.method == "POST" and request.headers.get("x-cc2-request") != "1":
            return JSONResponse({"error": "marker"}, 403)
        return await call_next(request)

    def form(body: bytes) -> dict:
        text = body.decode("utf-8").rstrip("\r\n")
        out = {}
        for line in text.split("\n"):
            if "=" not in line:
                raise HTTPException(400, "Invalid form")
            k, v = line.split("=", 1)
            if not k or k in out or len(k) >= 16 or len(v) > 96:
                raise HTTPException(400, "Invalid form")
            out[k] = v
        return out

    @app.get("/api/spools")
    def spools():
        with LOCK:
            STATE["requests"].append("GET /api/spools")
            return {"available": True, "enabled": P["enabled"], "slots": P["slots"], "spools": P["spools"],
                    "job": P["job"], "log": []}

    @app.get("/api/canvas")
    def canvas():
        return {"telemetry": {"result": {"canvas_info": {"canvas_list": [{"tray_list": [
            {"tray_id": s["slot"], "status": s["printer"]["status"]} for s in P["slots"][:4]]}]}}}}

    @app.post("/api/spools/enable")
    async def enable(request: Request):
        body = await request.body()
        if body not in (b"on", b"off"):
            raise HTTPException(400, "Invalid switch")
        with LOCK:
            P["enabled"] = body == b"on"
            STATE["requests"].append(f"POST enable {body.decode()}")
        return {"ok": True, "enabled": P["enabled"]}

    @app.post("/api/spools/save")
    async def save(request: Request):
        f = form(await request.body())
        with LOCK:
            STATE["requests"].append("POST save " + json.dumps(f, ensure_ascii=False))
            if f.get("id"):
                sp = next((s for s in P["spools"] if s["id"] == f["id"]), None)
                if not sp:
                    return JSONResponse({"ok": False, "error": "Unknown spool"}, 404)
                if "remaining" in f:
                    return JSONResponse({"ok": False, "error": "Correct the remaining weight with a weigh-in"}, 400)
                for k in ("name", "brand", "material", "color", "note"):
                    if k in f:
                        sp[k] = f[k]
                if "archived" in f:
                    sp["archived"] = f["archived"] == "1"
                return {"ok": True, "id": sp["id"]}
            if not f.get("material") or not f.get("color"):
                return JSONResponse({"ok": False, "error": "A spool needs a material and a colour"}, 400)
            sp = {"id": next(P["ids"]), "name": f.get("name", ""), "brand": f.get("brand", ""), "material": f["material"],
                  "color": f["color"], "diameter": float(f.get("diameter", 1.75)), "density": float(f.get("density", 1.24)),
                  "net": float(f.get("net", 1000)), "remaining": float(f.get("remaining", f.get("net", 1000))),
                  "tare": float(f.get("tare", 0)), "low": 0.0, "price": float(f.get("price", 0)),
                  "note": f.get("note", ""), "created": 0, "used": 0, "archived": False}
            P["spools"].append(sp)
            return JSONResponse({"ok": True, "id": sp["id"]}, 201)

    @app.post("/api/spools/assign")
    async def assign(request: Request):
        f = form(await request.body())
        with LOCK:
            STATE["requests"].append("POST assign " + json.dumps(f))
            slot = int(f["slot"])
            for s in P["slots"]:
                if f["spool"] and s["spool"] == f["spool"]:
                    s["spool"] = ""
            P["slots"][slot]["spool"] = f["spool"]
            P["slots"][slot]["question"] = ""
        return {"ok": True}

    return app


# ---------- Spoolman ----------

def make_spoolman() -> FastAPI:
    app = FastAPI()

    def get(sid: int):
        if sid not in STATE["spoolman"]:
            raise HTTPException(404, "not found")
        return STATE["spoolman"][sid]

    @app.get("/api/v1/spool")
    def spools(allow_archived: bool = False):
        return [s for s in STATE["spoolman"].values() if allow_archived or not s["archived"]]

    @app.get("/api/v1/spool/{sid}")
    def spool(sid: int):
        return get(sid)

    @app.patch("/api/v1/spool/{sid}")
    async def patch(sid: int, request: Request):
        body = await request.json()
        with LOCK:
            sp = get(sid)
            for key, value in (body.get("extra") or {}).items():
                if value is not None and key not in STATE["fields"]:
                    raise HTTPException(400, f"Unknown extra field {key}")
            sp.update({k: v for k, v in body.items() if k in ("location",)})
            for key, value in (body.get("extra") or {}).items():
                if value is None:
                    sp["extra"].pop(key, None)
                else:
                    sp["extra"][key] = value
            STATE["requests"].append(f"PATCH spoolman {sid} {json.dumps(body, ensure_ascii=False)}")
        return sp

    @app.get("/api/v1/field/spool")
    def fields():
        return [{"key": k, "entity_type": "spool", **v} for k, v in STATE["fields"].items()]

    @app.post("/api/v1/field/spool/{key}")
    async def add_field(key: str, request: Request):
        body = await request.json()
        with LOCK:
            STATE["fields"][key] = body
            STATE["requests"].append(f"POST field {key} {json.dumps(body, ensure_ascii=False)}")
        return fields()

    @app.put("/api/v1/spool/{sid}/use")
    async def use(sid: int, request: Request):
        body = await request.json()
        with LOCK:
            sp = get(sid)
            sp["used_weight"] += body["use_weight"]
            sp["remaining_weight"] -= body["use_weight"]
            STATE["requests"].append(f"PUT spoolman {sid} use {body['use_weight']}")
        return sp

    return app


# ---------- control ----------

def make_control() -> FastAPI:
    app = FastAPI()

    @app.get("/state")
    def state():
        return {**STATE, "printers": [{k: v for k, v in p.items() if k != "ids"} for p in PRINTERS]}

    @app.post("/insert")
    def insert(slot: int, spool: int, color: str = "0x000000", printer: int = 0):
        P = PRINTERS[printer]
        with LOCK:
            P["slots"][slot]["printer"] = {"status": 1}
            P["slots"][slot]["question"] = "inserted"
            P["slots"][slot]["spool"] = ""
            P["record"][f"channels{slot}"] = channel(10000 + spool, color=color)
        return {"ok": True}

    @app.post("/remove")
    def remove(slot: int, printer: int = 0):
        P = PRINTERS[printer]
        with LOCK:
            P["slots"][slot]["printer"] = {"status": 0}
            P["record"][f"channels{slot}"] = channel("")
        return {"ok": True}

    @app.post("/touch")
    def touch(slot: int, printer: int = 0):
        P = PRINTERS[printer]
        with LOCK:
            P["record"][f"channels{slot}"]["weight"] = "1000"
            P["slots"][slot]["spool"] = ""
            P["slots"][slot]["question"] = "changed"
        return {"ok": True}

    @app.post("/job")
    def job(active: int, printer: int = 0):
        P = PRINTERS[printer]
        with LOCK:
            P["job"]["active"] = bool(active)
        return {"ok": True}

    @app.post("/consume")
    def consume(slot: int, grams: float, printer: int = 0):
        P = PRINTERS[printer]
        with LOCK:
            sid = P["slots"][slot]["spool"]
            sp = next(s for s in P["spools"] if s["id"] == sid)
            sp["remaining"] -= grams
        return {"ok": True, "remaining": sp["remaining"]}

    @app.post("/authorize")
    async def authorize(request: Request):
        AUTHORIZED.add((await request.body()).decode().split()[1])
        return {"ok": True}

    return app


# ---------- SSH ----------

AUTHORIZED: set[str] = set()


class SSHServer(paramiko.ServerInterface):
    def __init__(self, P: dict):
        self.P = P

    def check_channel_request(self, kind, chanid):
        return paramiko.OPEN_SUCCEEDED if kind == "session" else paramiko.OPEN_FAILED_ADMINISTRATIVELY_PROHIBITED

    def get_allowed_auths(self, username):
        return "publickey"

    def check_auth_publickey(self, username, key):
        return paramiko.AUTH_SUCCESSFUL if username == "root" and key.get_base64() in AUTHORIZED else paramiko.AUTH_FAILED

    def check_channel_exec_request(self, channel_, command):
        def reply():  # forced command, whatever was sent
            with LOCK:
                data = json.dumps(self.P["record"]).encode()
            channel_.sendall(data)
            channel_.send_exit_status(0)
            channel_.close()
        threading.Timer(0.05, reply).start()  # after the exec request is acknowledged
        return True


def ssh_server(port: int, P: dict):
    host_key = paramiko.Ed25519Key.generate() if hasattr(paramiko.Ed25519Key, "generate") else None
    if host_key is None:
        from io import StringIO

        from cryptography.hazmat.primitives import serialization
        from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
        pem = Ed25519PrivateKey.generate().private_bytes(
            serialization.Encoding.PEM, serialization.PrivateFormat.OpenSSH, serialization.NoEncryption()).decode()
        host_key = paramiko.Ed25519Key.from_private_key(StringIO(pem))
    sock = socket.socket()
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind(("0.0.0.0", port))
    sock.listen(10)
    print(f"mock SSH on {port}, host key {host_key.fingerprint}", flush=True)
    while True:
        conn, _ = sock.accept()
        t = paramiko.Transport(conn)
        t.add_server_key(host_key)
        try:
            t.start_server(server=SSHServer(P))
        except Exception:
            pass


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--host", default="127.0.0.1", help="host part the CC2 Control guard expects")
    ap.add_argument("--printers", type=int, default=1)
    args = ap.parse_args()
    seed_spoolman()
    servers = [
        uvicorn.Server(uvicorn.Config(make_spoolman(), host="0.0.0.0", port=7912, log_level="warning")),
        uvicorn.Server(uvicorn.Config(make_control(), host="0.0.0.0", port=8099, log_level="warning")),
    ]
    for i in range(args.printers):
        P = printer_state()
        PRINTERS.append(P)
        threading.Thread(target=ssh_server, args=(2222 + i, P), daemon=True).start()
        servers.append(uvicorn.Server(uvicorn.Config(make_cc2(f"{args.host}:{8081 + i}", P), host="0.0.0.0",
                                                     port=8081 + i, log_level="warning")))
    threads = [threading.Thread(target=s.run, daemon=True) for s in servers[1:]]
    for t in threads:
        t.start()
    servers[0].run()


if __name__ == "__main__":
    main()