#!/usr/bin/env python

from pwn import *
import string
from parse import *
from hashlib import *
from itertools import chain, product

prompt1_template = "Please submit a printable string X, such that {}(X)[-6:] = {} and len(X) = {}"
vars_cd_template = "           c = {}, d = {};"
vars_ab_template = "           int a = {}, b = {};"
HASH = {"md5": md5, "sha1": sha1, "sha224": sha224, "sha256": sha256, "sha384": sha384, "sha512": sha512}

def pow(algo, hash, len):
    candidates = product(string.printable, repeat=int(len))
    for candidate in candidates:
        if HASH[algo](''.join(candidate)).hexdigest()[-6:] == hash:
            return ''.join(candidate)

def find_solution_strings(a, b, c, d):
    solution_string = "++--"
    for c_diff in range(-509, 509):
        for d_diff in range(-509, 509):
            if a*(c + c_diff) + b*(d + d_diff) == 1:
                print("Solution found!")
                solution_string_c = abs(c_diff)*"--" if c_diff < 0 else c_diff*"++"
                solution_string_d = abs(d_diff)*"--" if d_diff < 0 else d_diff*"++"
                return solution_string + solution_string_c, solution_string + solution_string_d

def exploit():
    r = remote("76.74.177.238", "7777")

    print("##### Step 1: Proof of work #####")
    prompt1_request = r.recvline()
    prompt1_data = parse(prompt1_template, prompt1_request)
    print("Attempting PoW: {}(X) = {}, len(X) = {}".format(prompt1_data[0], prompt1_data[1], prompt1_data[2]))
    correct_X = pow(prompt1_data[0], prompt1_data[1], prompt1_data[2])
    print("Found PoW: {}".format(correct_X))
    r.sendline(correct_X)

    print("##### Step 2: Find equation solution #####")
    r.recvline_endswith("[Q]uit!")
    r.sendline("g")
    r.recvline_endswith("d = 0;")
    vars_cd_line = r.recvline()
    vars_ab_line = r.recvline()
    vars_cd_data = parse(vars_cd_template, vars_cd_line)
    vars_ab_data = parse(vars_ab_template, vars_ab_line)
    a = int(vars_ab_data[0])
    b = int(vars_ab_data[1])
    c = int(vars_cd_data[0])
    d = int(vars_cd_data[1])
    print("Challenge: a = {}, b = {}, c = {}, d = {}".format(a, b, c, d))
    solution_c, solution_d = find_solution_strings(a, b, c, d)
    r.recvline_endswith("[Q]uit!")
    r.sendline("t")
    r.recvline_endswith("first variable:")
    r.sendline(solution_c)
    r.recvline_endswith("second variable:")
    r.sendline(solution_d)
    flag_line = r.recvline()
    r.close()
    print(flag_line)

if __name__ == "__main__":
    exploit()