[ CCE 2021] ptmd

#!python3
import socketserver
import random
import os,sys

welcome_msg = b"""----------------------------------------------------------------------
This is calc challenge.
Analyze the given equation and fill in the blanks.
Flag is awarded after passing 100 stages.
You can use +, -, *, /
----------------------------------------------------------------------
"""

op=['+','-','*','/']

def calculator(a,b,o):
    if o=='+':
        return a+b
    elif o=='-':
        return a-b
    elif o=='*':
        return a*b
    else:
        return a/b

def make_equation():
    a=random.randint(1, 0x1000)
    b=random.randint(1, 0x1000)
    op_=op[random.randint(0,3)]
    return a,b,op_

def win():
    result = ''
    with open('/flag', 'r') as f:
        result += f.read()
    return result.encode()

class MyTCPHandler(socketserver.BaseRequestHandler):
    def handle(self):

        s = self.request
        s.settimeout(10)
        s.send(welcome_msg)

        success=1
        for i in range(100):
            s.send("stage {}\n".format(str(i)).encode())
            a,b,op_=make_equation()
            c=calculator(a,b,op_)
            s.send('{} ? {} = {}\n>>'.format(str(a),str(b),str(c)).encode())
            op_=s.recv(1).decode()
            c_=calculator(a,b,op_)
            if(c_!=c):
                s.send(b"fail")
                s.close()
                success=0
                break
        if(success):
            s.send(win())
            s.close()

if __name__ == "__main__":

    if len(sys.argv) != 2:
        print("[-] usage : {} port".format(sys.argv[0]))
        sys.exit(-1)

    HOST, PORT = "0.0.0.0", int(sys.argv[1])

    try:
        socketserver.ThreadingTCPServer.allow_reuse_address = True
        server = socketserver.ThreadingTCPServer((HOST, PORT), MyTCPHandler)
        server.serve_forever()
    except:
        server.server_close()

문제 코드는 위와 같다. 서버에 접속하면 총 100개의 stage가 주어지고 해당 계산에 맞는 수학 기호를 입력하면 문제가 풀린다. 하지만 직접 입력하면 엔터도 같이 들어가기 때문에 파이썬 send함수를 사용해서 보내야한다. 코드는 아래와 같다.

from pwn import *
from parse import *

#context.log_level='debug'
r=remote("20.194.123.97",11111)

for i in range(100):
    r.recvuntil(b"stage ")
    r.recvuntil(b"\n")
    hello=r.recvuntil(b">>")
    print(b"[+] " + hello)
    result=hello.decode('utf-8')
    print(result)
    result=result.replace('\n',' ')
    result2=result.split(' ')
    if(float(result2[0])-float(result2[2])==float(result2[4])):
        r.send(b"-")
    if(float(result2[0])+float(result2[2])==float(result2[4])):
        r.send(b"+")
    if(float(result2[0])*float(result2[2])==float(result2[4])):
        r.send(b"*")
    if(float(result2[0])/float(result2[2])==float(result2[4])):
        r.send(b"/")
    result=""
    result2=[]

r.interactive()