Solver for you wa shockwave, PlaidCTF 2020

solver.py

from z3 import *

check_data = [
        [2, 5, 12, 19, 3749774],
        [2, 9, 12, 17, 694990],
        [1, 3, 4, 13, 5764],
        [5, 7, 11, 12, 299886],
        [4, 5, 13, 14, 5713094],
        [0, 6, 8, 14, 430088],
        [7, 9, 10, 17, 3676754],
        [0, 11, 16, 17, 7288576],
        [5, 9, 10, 12, 5569582],
        [7, 12, 14, 20, 7883270],
        [0, 2, 6, 18, 5277110],
        [3, 8, 12, 14, 437608],
        [4, 7, 12, 16, 3184334],
        [3, 12, 13, 20, 2821934],
        [3, 5, 14, 16, 5306888],
        [4, 13, 16, 18, 5634450],
        [11, 14, 17, 18, 6221894],
        [1, 4, 9, 18, 5290664],
        [2, 9, 13, 15, 6404568],
        [2, 5, 9, 12, 3390622],
]

def zz_helper(x, y, z):
    if y > z:
        # [1 (z - x)]
        return [1, (z - x)]

    # [y (x + y) z]
    c = zz_helper(y, x + y, z)
    a = c[0]
    b = c[1]

    # [b x]
    if b >= x:
        # [(2*a + 1) (b - x)
        return [(2*a)+1, (b - x)]

    # [(2*a) + 0, b]
    return [(2*a), b]


def zz(x):
    return zz_helper(1, 1, x)[0]

if __name__ == '__main__':
    flag = [BitVec("flag_{}".format(i), 32) for i in range(42)]
    checksum = BitVec("checksum", 32)
    s = Solver()

    for f in flag:
        s.add(And(f > 0, f <= 127))

    zz_map = Array('Map', BitVecSort(32), BitVecSort(32))

    for i in range(33000):
        s.add(zz_map[i] == zz(i))

    i = 1

    while i <= 21:
        index_1 = (2 * i) - 1
        index_2 = (2 * i)

        v1 = flag[index_1 - 1] * 256
        v2 = flag[index_2 - 1]

        checksum ^= zz_map[v1 + v2]
        i += 1

    s.add(checksum == 5803878)

    for check in check_data:
        i, j, k, l, target = check

        # Mysterious call to call-local in between computations
        i_sum = zz_map[(flag[2*i] << 8) | flag[(2*i) + 1]]
        j_sum = zz_map[(flag[2*j] << 8) | flag[(2*j) + 1]]
        k_sum = zz_map[(flag[2*k] << 8) | flag[(2*k) + 1]]
        l_sum = zz_map[(flag[2*l] << 8) | flag[(2*l) + 1]]

        s.add((i_sum ^ j_sum ^ k_sum ^ l_sum) == BitVecVal(target , 32))

    s.add(flag[0] == ord('P'))
    s.add(flag[1] == ord('C'))
    s.add(flag[2] == ord('T'))
    s.add(flag[3] == ord('F'))
    s.add(flag[4] == ord('{'))
    s.add(flag[-1] == ord('}'))

    print(s.check())
    m = s.model()

    res = ""

    for f in flag:
        if m[f] is not None:
            res += chr(m[f].as_long())

    print(res)