From e4b65d544f7fe412816d69c36e90d36d5cfa7a96 Mon Sep 17 00:00:00 2001 From: Michael <20937441+KMikeeU@users.noreply.github.com> Date: Fri, 8 May 2026 00:18:28 +0200 Subject: [PATCH] feat: code generation --- .env.example | 2 + Makefile | 23 +++++++- codegen/.python-version | 1 + codegen/README.md | 0 codegen/generate.py | 113 ++++++++++++++++++++++++++++++++++++++++ codegen/pyproject.toml | 7 +++ 6 files changed, 144 insertions(+), 2 deletions(-) create mode 100644 .env.example create mode 100644 codegen/.python-version create mode 100644 codegen/README.md create mode 100644 codegen/generate.py create mode 100644 codegen/pyproject.toml diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..cdde6a6 --- /dev/null +++ b/.env.example @@ -0,0 +1,2 @@ +FLAG_LEN=32 +FLAG_PREFIX=ctf diff --git a/Makefile b/Makefile index 78b91ed..a6dae17 100644 --- a/Makefile +++ b/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 $@ diff --git a/codegen/.python-version b/codegen/.python-version new file mode 100644 index 0000000..24ee5b1 --- /dev/null +++ b/codegen/.python-version @@ -0,0 +1 @@ +3.13 diff --git a/codegen/README.md b/codegen/README.md new file mode 100644 index 0000000..e69de29 diff --git a/codegen/generate.py b/codegen/generate.py new file mode 100644 index 0000000..1ae2840 --- /dev/null +++ b/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]} ") + 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 + +#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() diff --git a/codegen/pyproject.toml b/codegen/pyproject.toml new file mode 100644 index 0000000..115b41c --- /dev/null +++ b/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 = []