# UMASS{th3_spi11ag3_byt3c0d3_ctftim3_hack}

import types
import pickle
import random
import string

from flags import FLAGS


# Code type stuff

def serialize_function(f):
    return pickle.dumps([
        f.__code__.co_argcount,
        f.__code__.co_posonlyargcount,
        f.__code__.co_kwonlyargcount,
        f.__code__.co_nlocals,
        f.__code__.co_stacksize,
        f.__code__.co_flags,
        f.__code__.co_code,
        f.__code__.co_consts,
        f.__code__.co_names,
        tuple([var_name() for _ in f.__code__.co_varnames]),
        "<OwOwhatsthis>",
        fn_name(),# f.__code__.co_name,
        42069666,
        f.__code__.co_lnotab,
        f.__code__.co_freevars,
        f.__code__.co_cellvars
])

def deserialize_function(encoded):
    return types.FunctionType(types.CodeType(*pickle.loads(encoded)), globals())


# Symbolic obfuscation. This code for generating flags actually isn't used; I
# used the REPL to generate my thousand flags and hard-coded it into `flags.py`.

# >>> gen_flag()
# 'UMASS{th3_spi11ag3_byt3c0d3_ctftim3_hack}'
# >>> >>> gen_flag()
# 'UMASS{an_umass_wat3r_music_sarniqu3_m3m0ry_babushka}'
# >>> gen_flag()
# 'UMASS{a1r3ady_digging_scan_pick13_m3m0ry_buff3r_gam3}'

dic = [
    "python",
    "babushka",
    "pickle",
    "flag",
    "umass",
    "hack",
    "hacker",
    "pwn",
    "oi",
    "the",
    "a",
    "time",
    "already",
    "quickly",
    "lizard",
    "sarnique",
    "c",
    "overflow",
    "buffer",
    "rop",
    "gadget",
    "reversing",
    "flag",
    "key",
    "eleet",
    "spillage",
    "listerine",
    "clean",
    "bytecode",
    "nop",
    "ram",
    "memory",
    "gdb",
    "scan",
    "game",
    "cow",
    "dog",
    "fobrincate",
    "foo",
    "digging",
    "owo",
    "ctftime",
    "water",
    "hydrate",
    "jabber",
    "music",
    "rock",
    "an",
    "your",
    "my",
]

def elite_speak(word):
    leet_word = word.replace("e","3")
    leet_word = leet_word.replace("E","3")
    leet_word = leet_word.replace("l","1")
    leet_word = leet_word.replace("L","1")
    leet_word = leet_word.replace("o","0")
    leet_word = leet_word.replace("O","0")
    return leet_word

def gen_flag():
    return "UMASS{" + elite_speak("_".join(random.choices(dic, k=random.randrange(4, 8)))) + "}"

def fn_name():
    return ''.join(random.choices(string.ascii_uppercase, k=20))

def var_name():
    return ''.join(random.choices(string.ascii_uppercase, k=10))


# Code obfuscation. It isn't particularly sophisticated. Basically, we
#   - Generate some "conditions", which involve some bitwise comparison of parts
#     of the input.
#   - Randomize the order.
#   - Randomly intersperse code that does nothing.

def make_xor_cond(n, c):
    c = ord(c)
    up = c & 0b11110000
    lo = c & 0b00001111
    return "    b = b and (((ord(s[{n}]) & 0xf0) ^ {up} == 0) and ((ord(s[{n}]) & 0x0f) ^ {lo} == 0))".format(n = n, up = up, lo = lo)

def make_equals_cond(n, c):
    c = ord(c)
    return "    b = b and (ord(s[{n}]) == {c})".format(n = n, c = c)

def make_shift_cond(n, c):
    c = ord(c)
    return "    b = b and ((((ord(s[{n}]) & 0b10000000) >> 7) == {a}) and (((ord(s[{n}]) & 0b01000000) >> 6) == {b}) and (((ord(s[{n}]) & 0b00100000) >> 5) == {c}) and (((ord(s[{n}]) & 0b00010000) >> 4) == {d}) and (((ord(s[{n}]) & 0b00001000) >> 3) == {e}) and (((ord(s[{n}]) & 0b00000100) >> 2) == {f}) and (((ord(s[{n}]) & 0b00000010) >> 1) == {g}) and (((ord(s[{n}]) & 0b00000001) >> 0) == {h}))".format(
        n = n,
        a = (c & 0b10000000) >> 7,
        b = (c & 0b01000000) >> 6,
        c = (c & 0b00100000) >> 5,
        d = (c & 0b00010000) >> 4,
        e = (c & 0b00001000) >> 3,
        f = (c & 0b00000100) >> 2,
        g = (c & 0b00000010) >> 1,
        h = (c & 0b00000001) >> 0,
    )

def random_cond(n, c):
    i = random.randrange(0, 3)
    if i == 0:
        return make_xor_cond(n, c)
    elif i == 1:
        return make_equals_cond(n, c)
    elif i == 2:
        return make_shift_cond(n, c)

def benign_action():
    i = random.randrange(0, 6)
    if i == 0:
        return "    i = {}".format(random.randint(0, 69))
    if i == 1:
        return "    c = ord('{}') ^ ord(s[0])".format(random.choice(string.ascii_letters))
    if i == 2:
        return "    if ord(s[0]) & 0xf0 == {}:\n        i = {}\n    else:\n        i = {}".format(random.randint(0, 47), random.randint(0, 69), random.randint(0, 69))
    if i == 3:
        return "    copy.copy(s)"
    if i == 4:
        return "    i = len(s)"
    if i == 5:
        return "    if s[:5] != \"UMASS\":\n        return [False]"

FUNCTIONS = []
ENTRYFN = None

def gen_fn(i):
    global ENTRYFN
    flag = FLAGS[i - 1]
    conditions = [random_cond(i, c) for (i, c) in enumerate(flag)]
    end = "    b = [b]"
    if i > 1 and i % 4 == 1:
        end += "\n    b = {}(s) + b".format(str(FUNCTIONS[i - 2]))
    elif i > 1:
        end += "\n    t = types.FunctionType(types.CodeType(*pickle.loads({})), globals())\n    b = t(s) + b".format(str(FUNCTIONS[i - 2]))
    combiner = "    return b"
    if i == 25:
        combiner = "    combiner = {}\n    combiner = types.FunctionType(types.CodeType(*pickle.loads(combiner)), globals())\n    return [combiner(b)]".format(str(serialize_function(exactlyfourteenth_combiner)))
    elif i % 25 == 0:
        combiner = "    combiner = {}\n    combiner = types.FunctionType(types.CodeType(*pickle.loads(combiner)), globals())\n    return [combiner(b)]".format(str(serialize_function(any_combiner)))
    realcode = conditions + [benign_action() for _ in range(random.randrange(0, 13))]
    fname = fn_name()
    source_code = """
def {fname}(s):
    b = True
{code}
{end}
{combiner}
""".format(
    fname = fname,
    code = "\n".join(random.sample(realcode, len(realcode))),
    end = end,
    combiner = combiner
)
    try:
        exec(source_code)
    except Exception:
        print(source_code)
        raise "shit"
    if i % 4 == 1:
        print(source_code)
        ENTRYFN = fname
    if i % 4 == 0:
        print(source_code)
        FUNCTIONS.append(fname)
    else:
        FUNCTIONS.append(serialize_function(eval(fname)))
    return serialize_function(eval(fname))

# n.b. this is renamed to any_combiner in the pickled constants
def exactlyfourteenth_combiner(bools):
    if len(bools) < 14:
        return False
    bools[13] = False if bools[13] else True
    code = b'\x80\x04\x95\xbd\x00\x00\x00\x00\x00\x00\x00]\x94(K\x01K\x00K\x00K\x04K\x03KCC4d\x01}\x01|\x00D\x00]\x0c}\x02|\x01o\x12|\x02}\x01q\x08d\x02}\x03|\x00D\x00]\x0c}\x02|\x03|\x02O\x00}\x03q\x1e|\x01p2|\x03S\x00\x94N\x88\x89\x87\x94)(\x8c\nHMVDMTSAVI\x94\x8c\nBJWZRDQNFU\x94\x8c\nPPTMEAHGII\x94\x8c\nYRIMPWUSDG\x94t\x94\x8c\x0e<OwOwhatsthis>\x94\x8c\x0cany_combiner\x94J\xa2\xee\x81\x02C\x0e\x00\x01\x04\x01\x08\x01\n\x01\x04\x01\x08\x01\n\x01\x94))e.'
    combiner = types.FunctionType(types.CodeType(*pickle.loads(code)), globals())
    return not combiner(bools)

# n.b. only used in the constant code above
def any_combiner(bools):
    b = True
    for bp in bools:
        b = b and bp
    c = False
    for bp in bools:
        c |= bp
    return (b or c)

def any_combiner(bools):
    b = True
    for bp in bools:
        b = b and bp
    c = False
    for bp in bools:
        c |= bp
    return ((b or c) and False) or bools[0]

if __name__ == "__main__":
    print("import copy")
    print("import pickle")
    print("import types")
    for i in range(1, 501):
        gen_fn(i)
    print("""
def check_key(s):
    return {}(s.ljust(128))

if __name__ == "__main__":
    key = input("Oi, babushka, what's the key? ")
    combiner = types.FunctionType(types.CodeType(*pickle.loads(b'\\x80\\x04\\x95\\xcc\\x00\\x00\\x00\\x00\\x00\\x00\\x00]\\x94(K\\x01K\\x00K\\x00K\\x04K\\x03KCC@d\\x01}}\\x01|\\x00D\\x00]\\x0c}}\\x02|\\x01o\\x12|\\x02}}\\x01q\\x08d\\x02}}\\x03|\\x00D\\x00]\\x0c}}\\x02|\\x03|\\x02O\\x00}}\\x03q\\x1e|\\x01s4|\\x03r8d\\x02p>|\\x00d\\x03\\x19\\x00S\\x00\\x94(N\\x88\\x89K\\x00t\\x94)(\\x8c\\nOUVCHXMRZO\\x94\\x8c\\nGDOZHYKENT\\x94\\x8c\\nTZCLOUEGVM\\x94\\x8c\\nMMBJDKSFLR\\x94t\\x94\\x8c\\x0e<OwOwhatsthis>\\x94\\x8c\\x0cany_combiner\\x94J\\xa2\\xee\\x81\\x02C\\x0e\\x00\\x01\\x04\\x01\\x08\\x01\\n\\x01\\x04\\x01\\x08\\x01\\n\\x01\\x94))e.')), globals())

    if combiner(check_key(key)):
        print("Yes!")
    else:
        print("No!")
    """.format(ENTRYFN))