rctf-numbers-solver.py

#!/usr/bin/python

import pwn
import hashlib
import re
pwn.context.log_level = 'error'

import sys
import itertools
import collections

Feedback = collections.namedtuple('Feedback', ['correct', 'close'])

def generate_initial_pool(choices, holes):
  '''Generates the initial set of possible answers.'''
  return list(itertools.permutations(range(choices),holes))
  # return list(itertools.permutations(*[range(choices) for _ in xrange(holes)]))

def find_correct(actual, guess):
  '''Finds the sum of all correct matches.'''
  return sum([1 for (a, b) in zip(actual, guess) if a == b])


def remove_correct(actual, guess):
  '''Removes all correct matches from two "rows"'''
  actual2 = [a for (a, b) in zip(actual, guess) if a != b]
  guess2 = [b for (a, b) in zip(actual, guess) if a != b]
  return actual2, guess2


def find_close(actual, guess):
  '''Finds the sum of all close matches.'''
  actual, guess = remove_correct(actual, guess)

  close = 0
  for possible in guess:
    if possible in actual:
      del actual[actual.index(possible)]
      close += 1
  return close


def get_feedback(actual, guess):
  '''Compares two "rows" to each other and returns feedback.'''
  return Feedback(find_correct(actual, guess), find_close(actual, guess))


def is_match(guess, feedback, possible):
  '''Returns true if hypothetical could be the answer given the feedback
  and the guess'''
  return feedback == get_feedback(possible, guess)


def filter_pool(pool, guess, feedback):
  '''Filters through the pool of possibilities and removes ones which
  couldn't possibly be the answer.'''
  for possible in pool:
    if is_match(guess, feedback, possible) and (possible != guess):
      yield possible


def make_guess(pool, feedback):
  '''Makes an educated guess between the pool of possibilities and
  the user feedback.'''
  min_length = float('infinity')
  best_choice = None
  for possible in pool:
    length = len(list(filter_pool(pool, possible, feedback)))
    if min_length > length:
      min_length = length
      best_choice = possible
  return best_choice

# nc 149.28.139.172 10002

powChars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"

p = pwn.remote("149.28.139.172",10002)
data = p.recv()
[part1,hashTarget] = data.rstrip().split(" == ")
hashSeed = part1.split("+")[1].rstrip(")")

def dopow(seed,target):
  print "Calculating target hash %s for seed %s" % (target,seed)
  for c1 in powChars:
    for c2 in powChars:
      for c3 in powChars:
        for c4 in powChars:
          s = "%c%c%c%c" % (c1,c2,c3,c4)
          if hashlib.sha256(s + seed).hexdigest() == target:
            return s

pow = dopow(hashSeed,hashTarget)
p.recv()
print "PoW Complete: %s" % pow
p.sendline(pow)

# Give me 4 numbers, in[0, 10), You can only try 6 times
# $ 1 2 3 4
# Nope. 0, 1

def playGame(round):
  p.recvuntil("Give me ")
  numCount = int(p.recvuntil(" "),10)
  print "Got number count: %d" % numCount

  p.recvuntil("in[")
  minBound = int(p.recvuntil(",").rstrip(","),10)
  print "Got minimum bound: %d" % minBound

  p.recvuntil(" ")
  maxBound = int(p.recvuntil(")").rstrip(")"),10)
  print "Got maximum bound: %d" % maxBound

  # flush buffer then go
  # print p.recvall(timeout=0.5)
  print "OK, Playing game %d numbers between [%d:%d]" % (numCount,minBound,maxBound)
  numbers = range(minBound,maxBound)
  holes = numCount

  # flush buffer...
  p.recv()

  pool = generate_initial_pool(maxBound - minBound, numCount)
  guess = [minBound + i for i in range(0,numCount)]
  # guess = [0 if (i < (holes / 2)) else 1 for i in range(holes)]
  while True:
    actualGuess = [str(minBound + x) for x in guess]
    print "Sending: %s" % " ".join(actualGuess)
    p.sendline(" ".join(actualGuess))
    data = p.recv()
    print data
    if "Wrong" in data:
      correct = 0
      close = 0
    elif "Nope." in data:
      (correct, close) = data.rstrip().lstrip("Nope. ").split(", ")
      correct = int(correct)
      close = int(close)
    else:
      if round == 8:
        print p.recv()
      else:
        print p.recvline()
      return
    print "Result: %d correct, %d close" % (correct, close)
    feedback = Feedback(correct,close)
    pool = list(filter_pool(pool,guess,feedback))  
    guess = make_guess(pool,feedback)

r = 1
while True:
  playGame(r)
  r += 1

p.interactive()

p.close()