ångstromCTF 2021: Infinity Gauntlet

from pwn import *
from string import printable

def foo(param1=None, param2=None, result=None):
    if sum(1 if x is None else 0 for x in [param1, param2, result]) != 1:
        raise Exception('Something is wrong')
    if result is None:
        return str(int(param2 + 1 ^ param1 ^ 0x539))
    elif param1 is None:
        return str(int(result ^ (param2 + 1) ^ 0x539))
    elif param2 is None:
        return str(int((result ^ param1 ^ 0x539) - 1))

def bar(param1=None, param2=None, param3=None, result=None):
    if sum(1 if x is None else 0 for x in [param1, param2, param3, result]) != 1:
        raise Exception('Something is wrong')
    if result is None:
        return str(int((param3 + 1) * param2 + param1))
    elif param1 is None:
        return str(int(result - (param3 + 1) * param2))
    elif param2 is None:
        # can cause division by zero :-(
        return str(int((result - param1) / (param3 + 1)))
    elif param3 is None:
        # can cause division by zero :-(
        return str(int(((result - param1) / param2) - 1))

# context.log_level = 'DEBUG'
flag = [0 for i in range(264)]

conn = remote('shell.actf.co', 21700)
for i in range(1, 300):
    if i % 100 == 0:
        print('.', end='')
    conn.recvuntil(b'===\n')
    task = conn.recv(4)[:3]
    param1 = conn.recvuntil(b', ')[:-2]
    if task == b'foo':
        param2 = conn.recvuntil(b') = ')[:-4]
        result = conn.recvuntil(b'\n')[:-1]
    elif task == b'bar':
        param2 = conn.recvuntil(b', ')[:-2]
        param3 = conn.recvuntil(b') = ')[:-4]
        result = conn.recvuntil(b'\n')[:-1]
        param3 = None if param3 == b'?' else int(param3)
    param1 = None if param1 == b'?' else int(param1)
    param2 = None if param2 == b'?' else int(param2)
    result = None if result == b'?' else int(result)
    if task == b'foo':
        tmp = foo(param1, param2, result)
    elif task == b'bar':
        tmp = bar(param1, param2, param3, result)
    conn.sendline(tmp)
    if i > 50 and task == b'foo':
        idx = (int(tmp) >> 8) - i
        flag[idx] = int(tmp) % 256

conn.close()

flag_unxored = ''
for i, j in enumerate(flag):
    if j != 0:
        char_unxored = j ^ (i * 17) % 256
        if chr(char_unxored) not in printable and j > 128:
            char_unxored = (j - 128) ^ (i * 17) % 256
        flag_unxored += chr(char_unxored)
print('Flag: {}'.format(flag_unxored))