feat: code generation

.env.example

@@ -0,0 +1,2 @@
+FLAG_LEN=32
+FLAG_PREFIX=ctf

Makefile

@@ -1,3 +1,12 @@
+ifneq (,$(wildcard ./.env))
+  include .env
+  export
+else
+  $(warning .env file not found, loading defaults)
+  include .env.example
+  export
+endif
+
 RISCV_CC ?= riscv64-unknown-elf-gcc
 RISCV_ARCH_FLAGS ?= -march=rv32im -mabi=ilp32
 HOST_CC ?= gcc
@@ -8,6 +17,8 @@ EMU_SRC_FILES := $(wildcard src/emulated/*.c)
 EMU_OBJ_FILES := $(EMU_SRC_FILES:src/emulated/%.c=build/%.riscv.o)
 EMU_OBJ := build/emulated.riscv.elf
 
+GENERATED := build/include/generated.h
+
 HOST_SRC := src/main.c src/elf.c
 
 ifeq ($(OS),Windows_NT)
@@ -34,12 +45,20 @@ host: $(HOST_EXE)
 $(HOST_EXE): $(HOST_SRC) tinyriscv/src/tinyriscv.h src/config.h src/elf.h | build
 	$(HOST_CC) $(HOST_CFLAGS) $(HOST_LDFLAGS) -I src -I tinyriscv/src $(HOST_SRC) -o $@
 
-build/%.riscv.o: src/emulated/%.c | build
-	$(RISCV_CC) $(RISCV_ARCH_FLAGS) -c $< -o $@
+build/%.riscv.o: src/emulated/%.c $(GENERATED) | build
+	$(RISCV_CC) $(RISCV_ARCH_FLAGS) -I $(dir $(GENERATED)) -I src/emulated -c $< -o $@
 
 $(EMU_OBJ): $(EMU_OBJ_FILES) | build
 	$(RISCV_CC) $(RISCV_ARCH_FLAGS) -nostdlib -Wl,-e,entry -Wl,-Ttext=0 $(EMU_OBJ_FILES) -o $@
 
+
+.FORCE:
+
+$(GENERATED): .FORCE
+	mkdir -p $(dir $@)
+	python codegen/generate.py $(FLAG_PREFIX) $(FLAG_LEN) > $@
+
+
 build:
 	mkdir -p $@
 

codegen/.python-version

@@ -0,0 +1 @@
+3.13

codegen/generate.py

@@ -0,0 +1,113 @@
+import random
+import secrets
+import sys
+
+def eprint(*args, **kwargs):
+    print("[Codegen]", *args, file=sys.stderr, **kwargs)
+
+def custom_op_xor(value, amount):
+    # x[rd] = ((x[rs1] << (i32)imm) ^ 0xFF) >> (i32)imm;
+    return (((value << amount) ^ 0xFF) >> amount)
+
+
+manglers = [
+    (lambda value: custom_op_xor(value, 5), lambda value: f"trunc_xor({value}, 5)"),
+    (lambda value: custom_op_xor(value, 3), lambda value: f"trunc_xor({value}, 3)"),
+    (lambda value: custom_op_xor(value, 2), lambda value: f"trunc_xor({value}, 2)"),
+    (lambda value: custom_op_xor(value, 1), lambda value: f"trunc_xor({value}, 1)"),
+    (lambda value: value + 1, lambda value: f"{value} + 1"),
+    (lambda value: value + 128, lambda value: f"{value} + 128"),
+    (lambda value: value ^ 0xAB, lambda value: f"{value} ^ 0xAB"),
+    
+    # This is a trivial test mangler, do not use
+    # (lambda value: value, lambda value: f"{value}")
+]
+
+def mangle(value):
+    result = b""
+    c_manglers = []
+
+    for c in value:
+        mangler, c_mangler = random.choice(manglers)
+        # Mangle the input (str -> int)
+        # With overflow (mod 255)
+        result += (mangler((ord(c))) % 255).to_bytes()
+        c_manglers.append(c_mangler)
+
+    if len(result) != len(value):
+        eprint("WE GOT AN OVERFLOW")
+
+    return (result, c_manglers)
+
+def get_check_code(flag_part, user_part, c_mangler):
+    return f"""if(({flag_part}) != ({c_mangler(user_part)})){{
+    return false;
+}}
+    """
+
+def main():
+    if len(sys.argv) != 3:
+        eprint(f"Usage: {sys.argv[0]} <flag_prefix> <flag_length>")
+        exit(1)
+
+    # Flag prefix is not used, currently the flag is just used as-is
+    flag_prefix = sys.argv[1]
+    flag_len = int(sys.argv[2])
+
+    eprint(f"Generating flag with prefix {flag_prefix} and length {flag_len}")
+
+    # Generate a secret
+    secret = secrets.token_hex(flag_len)
+
+    # Change the bytes of the secret
+    mangled_secret, c_manglers = mangle(secret)
+
+    # Shuffle the secret, but keep track of the order
+    # We need this so we know the real location of the characters to compare
+    secret_shuffleorder = list(range(len(secret)))
+    random.shuffle(secret_shuffleorder)
+
+    # Shuffle the mangled secret according to the order
+    mangled_secret = b"".join([mangled_secret[x].to_bytes() for x in secret_shuffleorder])
+    # the mangler for secret[0] is no longer at c_manglers[0!]
+    c_manglers = [c_manglers[x] for x in secret_shuffleorder]
+
+
+    # Force binary representation
+    mangled_secret = "".join(f'\\x{b:02x}' for b in mangled_secret)
+
+    checks = []
+    for i, c_mangler in enumerate(c_manglers):
+        checks.append(get_check_code(f"MANGLED[{i}]", f"userflag[{secret_shuffleorder[i]}]", c_mangler))
+    
+    random.shuffle(checks)
+    check = "\n".join(checks)
+
+    # Just add one level of indent to make the generated code easier to read
+    check = check.replace("\n", "\n    ")
+
+    print(f"""// This file is automatically generated, do not modify!
+// The generated secret flag is "{secret}"
+
+#pragma once
+#include <stddef.h>
+
+#include "custom_instructions.h"
+
+static volatile uint8_t MANGLED[] = "{mangled_secret}";
+
+bool check(char* userflag, uint8_t len) {{
+    if (len != {len(secret)}) {{
+        return false;
+    }}
+
+    {check}
+
+    return true;
+}}
+
+""")
+
+
+if __name__ == "__main__":
+    main()

codegen/pyproject.toml

@@ -0,0 +1,7 @@
+[project]
+name = "codegen"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.13"
+dependencies = []