Red Team Village Defcon 2020 - Seeding Solution

seeding_solution.py

#!/usr/bin/env python3
# random.choice produces different values on Python 2.7 vs Python 3

from pwn import *
import random
from itertools import permutations

server_ip = '127.0.0.1'


def update_keyspace(key_space):

	r = remote(server_ip, 3000)

	data = r.recvuntil("The seed this run was:")
	seed = r.recvline().strip()
	#print("Seed: " + seed)

	data = r.recvuntil("The fourth iteration of this key is:")
	fourth_key = r.recvline().strip()
	#print("Fourth Key: " + fourth_key)

	r.close()

	key_set = set(fourth_key)
	key_space.update(key_set)

def get_seed_and_fourth_iter():

	r = remote(server_ip, 3000)

	data = r.recvuntil("The seed this run was:")
	seed = int(r.recvline().strip())
	
	data = r.recvuntil("The fourth iteration of this key is:")
	fourth_key = r.recvline().strip()
	
	r.close()

	return seed,fourth_key

def get_flag1(server_key):

	flag1 = ""
	flag2 = ""
	r = remote(server_ip, 3000)

	data = r.recvuntil("The seed this run was:")
	seed = r.recvline().strip()
	#print("Seed: " + seed)

	data = r.recvuntil("The fourth iteration of this key is:")
	fourth_key = r.recvline().strip()
	#print("Fourth Key: " + fourth_key)

	data = r.recvuntil("What is the key?: ")
	#print(data)

	# Send key
	r.sendline(server_key)

	# Get the first flag
	data = r.recvline().decode('utf-8')
	flag_str = "Heres the flag "
	idx = data.find(flag_str)
	if idx != -1:

		flag1 = data[idx + len(flag_str):].strip()
		#flag1 = r.recvline().strip()

		#Get the second flag
		val = get_nth_iteration(int(seed), server_key, 1337, len(fourth_key))
		
		# Send key
		r.send(val)

		# Get the first flag
		data = r.recvuntil("Heres the last")
		flag2 = r.recvline().strip()
	else:
		print(data)

	r.close()

	return flag1,flag2

def get_random_string(key, length): 
    result_str = ''.join(random.choice(key) for i in range(length)) 
    return result_str

def gen_ascii_strs(max_len):

	generator = permutations(string.ascii_letters, max_len)
	for p in generator:
		ascii_entry = ''.join(p)		
		return ascii_entry

def get_next_str(cur_pass, key_len):

	rand_str = get_random_string(cur_pass, key_len)
	    #print(rand_str)

	#Get idx of chars in string to resolve key
	fourth_str = rand_str 
	#print("Local Fourth: " + fourth_str)
	idx_arr = []
	for key_char in fourth_str:
		idx = cur_pass.index(key_char)
		idx_arr.append(idx)

	return idx_arr,fourth_str

def get_key_arr(idx_arr, guess_pass_len, seed_key):

	# print(idx_arr)
	orig_key = ['_'] * guess_pass_len
	j = 0
	for idx in idx_arr:
		orig_key[idx] = seed_key[j]
		j += 1

	return orig_key

def compare_keys(key1, key2):

	if len(key1) != len(key2):
		return False

	for i in range(len(key1)):
		if key1[i] != key2[i] and key1[i] != "_" and key2[i] != "_":
			return False

	return True

def get_key_space():

	key_space = set()
	for i in range(15):
		update_keyspace(key_space)
		key_space_len = len(key_space)
		#print("Unique chars: %d" % key_space_len)
		#rint(key_space)	
	return key_space

def get_candidate_key(key_len, seed, random_str):

	random_str_len = len(random_str)
	cur_pass = gen_ascii_strs(key_len)
	#print("Local Key: " + cur_pass)

	#Get index arr for the seed
	random.seed(seed) 
	for j in range(4):
		first_seed_idx_arr, fourth_str = get_next_str(cur_pass, random_str_len)
	
	# Get key arr for first sample server 4th
	ret_key = get_key_arr(first_seed_idx_arr, key_len, random_str)

	return ret_key


def get_candidate_key_length(key_space_len_min):

	#Get a sample seed and fourth iteration 
	first_seed,first_fourth_val = get_seed_and_fourth_iter()

	#Get another sample seed and fourth iteration 
	second_seed,second_fourth_val = get_seed_and_fourth_iter()

	#print("Sample1: Seed: " + str(first_seed))
	#print("Sample1: Fourth Key: " + first_fouth_val)
	#print("\n")
	#print("Sample2: Seed: " + str(second_seed))
	#print("Sample2: Fourth Key: " + second_fouth_val)
	#print("\n")

	random_str_len = len(second_fourth_val)
	#print("Random String Len: %d" % random_str_len)
	#print("\n")

	candidate_length = []
	for k in range(key_space_len_min, key_space_len_min + 10):

		guess_pass_len = k
		#print("[+] Key Len Guess:  %d" % guess_pass_len)

		# Get two keys with different seeds but the same length and compare them
		ret_key1 = get_candidate_key(guess_pass_len, first_seed, first_fourth_val)
		#print(ret_key1)
		ret_key2 = get_candidate_key(guess_pass_len, second_seed, second_fourth_val)
		#print(ret_key2)

		if compare_keys(ret_key1, ret_key2):
			candidate_length.append(guess_pass_len)

	#Possible key length
	return candidate_length

def get_key_length(key_space_len):

	#Get a candidate length
	length_set = None
	while length_set == None or len(length_set) > 1:

		candidate_key_len_list = get_candidate_key_length(key_space_len)
		temp_set = set(candidate_key_len_list)

		#print("[+] Set key:  \t\t" + str(temp_set))

		#If there's only one then return because we found it
		if(length_set) != None:
			length_set = length_set.intersection(temp_set)
		else:
			length_set = temp_set
			
		if len(length_set) == 1:
			break

	if len(length_set) == 0:
		print("[-] Unable to determine key length")
		sys.exit(1)

	#Key length
	return list(length_set)[0]

def get_server_key(key_len):

	final_key = ['_'] * key_len
	while True:
		cur_seed,fouth_iteration_str = get_seed_and_fourth_iter()
		ret_key1 = get_candidate_key(key_len, cur_seed, fouth_iteration_str)
		for i in range(key_len):
			if ret_key1[i] != "_":
				final_key[i] = chr(ret_key1[i])

		temp_key = ''.join(final_key)
		if "_" not in temp_key:
			break

	return ''.join(final_key)

def get_nth_iteration(seed, key, iter_num, random_str_len):
	#Get index arr for the seed
	random.seed(seed) 
	for j in range(iter_num):
		first_seed_idx_arr, fourth_str = get_next_str(key, random_str_len)
	return fourth_str

# Get the current key space
key_space = get_key_space()
key_space_len = len(key_space)
print("[+] Key space size: %d" % key_space_len)

# Get key length
key_len = get_key_length(key_space_len)
print("[+] Key length: %d" % key_len)

# Get key
server_key = get_server_key(key_len)
print("[+] Server key: %s" % server_key)

# Get the first flag
flag1,flag2 = get_flag1(server_key)
print("[+] %s" % flag1)
print("[+] %s" % flag2.decode('utf-8'))