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= (scanned roll in),
/remove?slot= (roll out, the record keeps its tag data), /putback?slot= (roll in without a scan),
/takeout?slot= (CC2 Control's "Take out"), /consume?slot=&grams=, /job?active=0|1,
/touch?slot= (touchscreen edit: weight falls back to 1000), /print?slot=&mm=&job=&result= (logged use of
a tray, charged to its spool if it has one), /weigh?spool=&remaining= (weigh-in in CC2 Control).
"""

from __future__ import annotations

import argparse
import itertools
import json
import socket
import sys
import threading
import time
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": []},
        "log": [],
        "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": [], "vendors": [], "filaments": []}
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 if vendor == "Sunlu" else 2 if vendor == "Elegoo" else 3, "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": P["log"]}

    @app.get("/api/printer")
    def printer_status():
        active = P["job"]["active"]
        return {"connected": True, "machine": {"status": 2 if active else 1, "progress": 42 if active else 0},
                "print": {"filename": P["job"].get("file", ""), "state": "printing" if active else "standby",
                          "current_layer": 23 if active else 0, "total_layers": 55, "remaining": 2940}}

    @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]
                for k in ("diameter", "density", "net", "tare", "low", "price"):
                    if k in f:
                        sp[k] = float(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/adjust")
    async def adjust(request: Request):
        f = form(await request.body())
        with LOCK:
            STATE["requests"].append("POST adjust " + json.dumps(f))
            sp = next((s for s in P["spools"] if s["id"] == f.get("id")), None)
            if not sp:
                return JSONResponse({"ok": False, "error": "Unknown spool"}, 404)
            value = float(f["remaining"])
            P["log"].insert(0, {"time": int(time.time()), "spool": sp["id"], "slot": -1, "kind": "correction",
                                "grams": value - sp["remaining"], "mm": 0.0, "job": "", "result": "set"})
            sp["remaining"] = value
        return {"ok": True, "remaining": value}

    @app.post("/api/spools/delete")
    async def delete(request: Request):
        f = form(await request.body())
        with LOCK:
            STATE["requests"].append("POST delete " + json.dumps(f))
            if not any(s["id"] == f.get("id") for s in P["spools"]):
                return JSONResponse({"ok": False, "error": "Unknown spool"}, 404)
            P["spools"][:] = [s for s in P["spools"] if s["id"] != f["id"]]
            for s in P["slots"]:
                if s["spool"] == f["id"]:
                    s["spool"] = ""
        return {"ok": True}

    @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", "lot_nr", "archived")})
            if "used_weight" in body:
                sp["used_weight"] = float(body["used_weight"])
                sp["remaining_weight"] = (sp.get("initial_weight") or 1000.0) - sp["used_weight"]
            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/vendor")
    def vendors():
        found = {s["filament"]["vendor"]["id"]: s["filament"]["vendor"] for s in STATE["spoolman"].values()}
        found.update({v["id"]: v for v in STATE["vendors"]})
        return list(found.values())

    @app.get("/api/v1/filament")
    def filaments():
        found = {s["filament"]["id"]: s["filament"] for s in STATE["spoolman"].values()}
        found.update({f["id"]: f for f in STATE["filaments"]})
        return list(found.values())

    @app.post("/api/v1/vendor")
    async def add_vendor(request: Request):
        body = await request.json()
        with LOCK:
            v = {"id": max([x["id"] for x in vendors()], default=0) + 1, **body}
            STATE["vendors"].append(v)
            STATE["requests"].append(f"POST vendor {json.dumps(body, ensure_ascii=False)}")
        return v

    @app.post("/api/v1/filament")
    async def add_filament(request: Request):
        body = await request.json()
        with LOCK:
            vendor = next(v for v in vendors() if v["id"] == body["vendor_id"])
            f = {"id": max([x["id"] for x in filaments()], default=0) + 1, "vendor": vendor, **body}
            STATE["filaments"].append(f)
            STATE["requests"].append(f"POST filament {json.dumps(body, ensure_ascii=False)}")
        return f

    @app.post("/api/v1/spool")
    async def add_spool(request: Request):
        body = await request.json()
        with LOCK:
            fil = next(f for f in filaments() if f["id"] == body["filament_id"])
            sid = max(STATE["spoolman"], default=0) + 1
            weight = body.get("initial_weight") or fil.get("weight") or 1000.0
            STATE["spoolman"][sid] = {"id": sid, "lot_nr": None, "remaining_weight": weight, "used_weight": 0.0,
                                      "initial_weight": weight, "spool_weight": body.get("spool_weight"),
                                      "location": body.get("location", ""), "archived": False, "price": body.get("price"),
                                      "extra": {}, "filament": fil}
            STATE["requests"].append(f"POST spool {json.dumps(body, ensure_ascii=False)}")
        return STATE["spoolman"][sid]

    @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}
            if P["slots"][slot]["spool"]:
                P["slots"][slot]["last"] = P["slots"][slot]["spool"]
            P["slots"][slot]["spool"] = ""
            P["slots"][slot]["question"] = ""
        return {"ok": True}

    @app.post("/putback")
    def putback(slot: int, printer: int = 0):
        P = PRINTERS[printer]
        with LOCK:
            P["slots"][slot]["printer"] = {"status": 1}
            P["slots"][slot]["spool"] = ""
            P["slots"][slot]["question"] = "inserted"
        return {"ok": True}

    @app.post("/takeout")
    def takeout(slot: int, printer: int = 0):
        P = PRINTERS[printer]
        with LOCK:
            P["slots"][slot]["last"] = P["slots"][slot]["spool"]
            P["slots"][slot]["spool"] = ""
            P["slots"][slot]["question"] = ""
        return {"ok": True}

    def log(P, kind, spool, slot, grams, mm, job="", result=""):
        P["log"].insert(0, {"time": int(time.time()), "spool": spool, "slot": slot, "kind": kind, "grams": grams,
                            "mm": mm, "job": job, "result": result})
        del P["log"][150:]

    @app.post("/print")
    def print_(slot: int, mm: float, job: str = "test.gcode", result: str = "complete", 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), None)
            d, rho = (sp["diameter"], sp["density"]) if sp else (1.75, 1.24)
            grams = mm * 3.141592653589793 * (d / 2) ** 2 * rho / 1000
            if sp:
                sp["remaining"] -= grams
            log(P, "print", sid, slot, -round(grams, 3), mm, job, result)
        return {"ok": True, "grams": grams}

    @app.post("/weigh")
    def weigh(spool: str, remaining: float, printer: int = 0):
        P = PRINTERS[printer]
        with LOCK:
            sp = next(s for s in P["spools"] if s["id"] == spool)
            log(P, "correction", spool, -1, remaining - sp["remaining"], 0, "", "weighed")
            sp["remaining"] = remaining
        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()