126 lines
3.8 KiB
Python
126 lines
3.8 KiB
Python
import random
|
|
import sys
|
|
|
|
def generate_flag(length):
|
|
characters = list("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")
|
|
return "".join([random.choice(characters) for _ in range(length)])
|
|
|
|
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):
|
|
variations = [
|
|
f"custom_jmp_neq({flag_part} ^ 0x7C, {c_mangler(user_part)}, label_false);",
|
|
f"custom_jmp_neq({flag_part}, {c_mangler(user_part)} ^ 0x7C, label_false);"
|
|
]
|
|
return random.choice(variations)
|
|
|
|
def main():
|
|
if len(sys.argv) != 3 and len(sys.argv) != 4:
|
|
eprint(f"Usage: {sys.argv[0]} <flag_format> <flag_length> [flag_seed]")
|
|
exit(1)
|
|
|
|
# Flag prefix is not used, currently the flag is just used as-is
|
|
flag_format = sys.argv[1]
|
|
flag_len = int(sys.argv[2])
|
|
|
|
if len(sys.argv) >= 4:
|
|
# We might have a seed
|
|
flag_seed = sys.argv[3]
|
|
if flag_seed.lower() != "false":
|
|
random.seed(flag_seed)
|
|
|
|
# Generate a secret
|
|
secret = generate_flag(flag_len)
|
|
secret = flag_format.replace("FLAG_GOES_HERE", secret)
|
|
|
|
eprint(f"Generated Flag: {secret}")
|
|
|
|
# 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) {{
|
|
custom_jmp_neq(len ^ 0x7C, {len(secret)}, label_false);
|
|
|
|
{check}
|
|
|
|
return true;
|
|
|
|
label_false:
|
|
return false;
|
|
}}
|
|
|
|
""")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|