Imaginary - ImaginaryCTF

#!/usr/bin/env python3

import random
from solve import solve

banner = '''
Welcome to the Imaginary challenge! I'm gonna give you 300 imaginary/complex number problems, and your job is to solve them all. Good luck!

Sample input: (55+42i) + (12+5i) - (124+15i)
Sample output: -57+32i

Sample input: (23+32i) + (3+500i) - (11+44i)
Sample output: 15+488i

(NOTE: DO NOT USE eval() ON THE CHALLENGE OUTPUT. TREAT THIS IS UNTRUSTED INPUT. Every once in a while the challenge will attempt to forkbomb your system if you are using eval(), so watch out!)
'''

flag = open("flag.txt", "r").read()
ops = ['+', '-']

print(banner)

for i in range(300):
	o = random.randint(0,50)
	if o > 0:
		nums = []
		chosen_ops = []
		for n in range(random.randint(2, i+2)):
			nums.append([random.randint(0,50), random.randint(0,50)])
			chosen_ops.append(random.choice(ops))
		out = ""
		for op, num in zip(chosen_ops, nums):
			out += f"({num[0]}+{num[1]}i) {op} "
		out = out[:-3]
		print(out)
		ans = input("> ")
		if ans.strip() == solve(out).strip():
			print("Correct!")
		else:
			print("That's incorrect. :(")
			exit()
	else:
		n = random.choice("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
		payload = f"__import__['os'].system('{n}(){{ {n}|{n} & }};{{{n}}}')"
		print(payload)
		input("> ")
		print("Correct!")

print("You did it! Here's your flag!")
print(flag)

'''
ictf{n1c3_y0u_c4n_4dd_4nd_subtr4ct!_49fd21bc}
'''

import socket
import re

HOST = "chal.imaginaryctf.org"
PORT = 42015
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

s.connect((HOST, PORT))

def process(s):
  lhs = []
  rhs = []

  lhs_r = r"(?<=\()\d+"
  lhs = re.findall(lhs_r, s)
  rhs_r = r".\d+(?=i)"
  rhs = re.findall(rhs_r, s)
  op_r = r" [+-] "
  op = re.findall(op_r, s)

  final_l = ""
  final_r = ""

  for i in range(len(op)):
    final_l += lhs[i] + op[i]
    final_r += rhs[i] + op[i]

  final_l += lhs[-1]
  final_r += rhs[-1]

  final_l = eval(final_l)
  final_r = eval(final_r)
  if final_r < 0:
    return f"{final_l}{final_r}i"
  else:
    return f"{final_l}+{final_r}i"

# First problem
prob = s.recv(1024).splitlines()[-2].decode("utf-8")
sol = bytes(process(prob) + "\n", "utf-8")
s.send(sol)

newline = bytes("\n", "utf-8")
i = 1
while i < 300:
  i += 1
  prob = s.recv(4096).decode("utf-8")
  if len(prob) in (1428, 2856):
    prob += s.recv(4096).decode("utf-8")
  print(prob, type(prob), len(prob), i)
  if "__" in prob:
    s.send(newline)
    continue
  sol = bytes(process(prob) + "\n", "utf-8")
  #print(prob)
  #print(sol, i)
  s.send(sol)

s.send(sol)

s.close()

ictf{n1c3_y0u_c4n_4dd_4nd_subtr4ct!_49fd21bc}