import random
import sys
import os

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)


hard_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}")
]

easy_manglers = [
    (lambda value: value + 1, lambda value: f"{value} + 1"),
    (lambda value: value - 1, lambda value: f"{value} - 1"),
    (lambda value: value * 2, lambda value: f"{value} * 2"),
    (lambda value: value ^ 0xAB, lambda value: f"({value} ^ 0xAB)"),
    (lambda value: value, lambda value: f"{value}")
]



def mangle(value, manglers):
    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, difficulty):

    if difficulty == "HARD":
        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);"
        ]
    else: 
        variations = [
            f"if({flag_part} != {c_mangler(user_part)}) {{ return 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("totallysecurepadding" + flag_seed + "sonoonecanguesstheflag!")

    difficulty = os.getenv("DIFFICULTY", "EASY")

    # 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, easy_manglers if difficulty == "EASY" else hard_manglers)

    # 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, difficulty))
    
    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) {{
    {f"custom_jmp_neq(len ^ 0x7C, {len(secret)}, label_false);" if difficulty == "HARD" else f"if (len != {len(secret)}) {{ return false; }}"}

    {check}

    return true;

    label_false:
        return false;
}}

""")


if __name__ == "__main__":
    main()