VolgaCTF mathproblem

volgactf_math.py

import telnetlib
import re
import time

flag = '{you_count_as_fast_as_a_calculator}'

class Telnet(telnetlib.Telnet):
    # inherit from Telnetlib and add new method
    def __init__(self,host,port):
        telnetlib.Telnet.__init__(self,host,port)
    # make easier when you want to send raw data to server
    def writeRawData(self,data):
        return self.get_socket().send(data)

    def recvRawData(self,size):
        return self.get_socket().recv(size)

# postfix handling
# DFS on equations
def rec(num, expr, on_stack):
    if ((len(num) == 0) and (on_stack == 1)):
        yield expr
        return
    choices = num[:]
    if (on_stack >= 2):
        choices += ['+', '-', '*', '/']
    for choice in choices:
        num_prime = num[:]
        on_stack_prime = on_stack
        if (choice in num):
            num_prime.remove(choice)
            on_stack_prime += 1
        else:
            on_stack_prime -= 1
        for ress in rec(num_prime, expr + [choice], on_stack_prime):
            yield ress

# evaluate a postfix expression
def evalExpr(expr):
    stack = []
    for i in expr:
        if (i in ['+', '-', '*', '/']):
            op2 = stack.pop()
            op1 = stack.pop()
            if   (i == '+'): stack.append(op1 + op2)
            elif (i == '-'): stack.append(op1 - op2)
            elif (i == '*'): stack.append(op1 * op2)
            elif (i == '/'):
                try:
                    stack.append(op1 // op2)
                except ZeroDivisionError:
                    return -6996699669966996699669966996L # so it's different from the result
        else:
            stack.append(i)
    return stack[0]

# convert an expression from postfix notation to infix notation
def toInfix(expr):
    res = []
    for i in expr:
        if (type(i) is int):
            res.append(i)
        else:
            op2 = res.pop()
            op1 = res.pop()
            res.append([op1, i, op2])
    return res

# convert an infix expression to string
def prettyPrint(infix):
    if (type(infix) is list):
        return '(' + ''.join(map(prettyPrint, infix)) + ')'
    return str(infix)

# connect
host = 'mathproblem.2015.volgactf.ru'
port = 8888
tel = Telnet(host, port)
# read title
s = tel.recvRawData(20480)
print s
# solve the task
for i in xrange(30):
    s = tel.recvRawData(20480)
    print s,
    nums = re.findall(r'Solve!\n([-\d\s]+)\ equals\ ([-\d]+)\n', s)[0]
    l = list(map(int, nums[0].split()))
    eq = int(nums[1])
    for ex in rec(l, list(), 0):
        val = evalExpr(ex)
        if (val == eq):
            res = prettyPrint(toInfix(ex))
            print res
            tel.writeRawData(res+'\n')
            break
# get flag
s = tel.recvRawData(20480)
print s
tel.close()