0ctf 2018 quals "udp" solution

solve.py

import sys
import pprint
import struct

TABLE_SIZE = 4000
table = [[]]
iterators = []
locks = set()

def request(cur=0, path_diff=2**64):
    global table, locks, iterators
    if cur == 1:                    # special case
        return path_diff            # accept all the shit
    locks.add(cur)
    for i in range(iterators[cur], TABLE_SIZE):
        if i == cur:
            continue
        if i in locks:
            continue
        if table[cur][i] == 0:
            continue
        res = request(i, min(table[cur][i], path_diff))
        if res:
            iterators[cur] = i      # save current position
            path_diff = res
            table[cur][i] -= path_diff
            table[i][cur] += path_diff
            if cur == 0 and table[cur][i] == 0:
                print(i)
            break
    else:
        path_diff = 0
    locks.remove(cur)
    return path_diff

def solve():
    total = 0
    while True:
        diff = request()
        if not diff:
            break
        total += diff
        if total % 100 == 0:
            print('total', total)
    print('\nfinal total =', total)
    print('flag{{{}}}'.format(hex(total)[2:]))

def main():
    global table, locks, iterators
    sys.setrecursionlimit(5000)
    with open('udp', 'rb') as f:
        data = f.read()
    start = data.find(struct.pack('<Q', 0x814BE171)) - 8
    data = data[start:start + TABLE_SIZE * TABLE_SIZE * 8]
    table = [list(struct.unpack('<{}Q'.format(TABLE_SIZE), data[i * TABLE_SIZE * 8:(i + 1) * TABLE_SIZE * 8])) for i in range(TABLE_SIZE)]
    iterators = [0 for i in range(TABLE_SIZE)]
    locks = set()
    solve()

if __name__ == '__main__':
    main()