Solver for YOU wa SHOCKWAVE (re 250) from PlaidCTF 2020

YOU_wa_SHOCKWAVE_solver.py

from struct import pack, unpack

# reversed from the challenge
def zz_helper(x, y, z):
    if y > z:
        return [1, z - x]
 
    a, b = zz_helper(y, x + y, z)
    if b >= x:
        return [2 * a + 1, b - x]
    return [2 * a, b]
 
def zz(x):
    return zz_helper(1, 1, x)[0]

# inverse transformation
def rev_zz(x):
    if x == 1:
        return 0
    x //= 2
    res = 0
    a = 1
    b = 1
    i = 0
    while (1 << i) <= x:
        if x & (1 << i):
            res += b
        a, b = b, a+b
        i += 1
    return res

for i in range(1000):
    assert rev_zz(zz(i)) == i
    assert zz(rev_zz(zz(i))) == zz(i)

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]
]

m = []
for (i, j, k, l, target) in check_data:
    row = [0] * 21 + [target]
    row[i] = 1
    row[j] = 1
    row[k] = 1
    row[l] = 1
    m.append(row)
m.append([1]*21 + [5803878])

for i in range(21):
    # select
    for j in range(i, 21):
        if m[j][i]:
            m[j], m[i] = m[i], m[j]
            break
    else:
        raise RuntimeError('fail')
    # clear others
    for j in range(21):
        if j != i and m[j][i]:
            for k in range(22):
                m[j][k] ^= m[i][k]

for row in m:
    print(row)

flag = b''
for i in range(21):
    r = rev_zz(m[i][-1])
    assert(zz(r) == m[i][-1])
    flag += pack('>H', r)
assert len(flag) == 42
print(flag.decode())