2016 ASIS CTF secuprim

import string
import sys
import re
from pwn import *
from itertools import product, permutations
from Crypto.Util.number import isPrime
from gmpy2 import isqrt
from multiprocessing import Process, Queue


def cnt_prime(a, b, output) :
	if a % 2 == 0 :
		a += 1
	cnt = 0
	for i in range(a, b, 2) :
		if isPrime(i) == True :
			cnt += 1
	output.put(cnt)

def cul(a, b) :
	output1 = Queue()
	output2 = Queue()
	procs = []
	start = a
	end = b
	mid = (a+b)/2

	procs.append(Process(target=cnt_prime, args=(start,mid,output1)))
	procs.append(Process(target=cnt_prime, args=(mid,end+1,output2)))

	for p in procs :
		p.start()

	cnt = output1.get() + output2.get()
	output1.close()
	output2.close()

	for p in procs :
		p.join()

	return cnt

if __name__ == '__main__' :
	maps = string.printable[:62]
	while(True) :
		s = remote('secuprim.asis-ctf.ir', 42738)
		p = log.progress('Working')

		p.status('Stage 1!')
		tmp = s.recvuntil('X: ').split('\n')[2].split('"')
		str_, hash_  = tmp[1], tmp[3][:-3]
		first = ''
		#for i in product(maps, repeat=4) :
		for i in permutations(maps, 4) :
			tmp = sha256sumhex(''.join(i)+str_)[:len(hash_)]
			if tmp == hash_ :
				first = ''.join(i)
				break
		s.sendline(first)
		s.recvuntil('---\n')

		for _ in range(2, 100) :
			p.status('Stage %s ing~'%(str(_)))
			try :
				tmp = s.recvuntil('like n such that')
				tmp = s.recvuntil('\n')
			except :
				tmp = s.recvall()
				print tmp
				if bool(re.search('ASIS',tmp)) :
					p.success('Clear !')
					sys.exit(1)
				break
			tmp = tmp.split(': ')[1].split()
			start, end = int(tmp[0]), int(tmp[4])
			cnt = cul(start, end)
			cnt += int(isqrt(end)-isqrt(start))
			s.sendline(str(cnt))
			log.success('Stage %s clear!'%(str(_)))