CTF101 - pygenere (multi key xor cipher)

ctf101_pygenere.py

from pwn import xor

def bytes_to_chunks(b, n_chunks):
	chunks = [[] for _ in range(n_chunks)]

	for i in range(len(b)):
		chunks[i % n_chunks].append(b[i])

	return [bytes(c) for c in chunks]

def good_keys(b):
	good = []
	for k in range(256):
		guess = xor(b, k)
		# Is only chars
		if min(guess.replace(b'\n', b'')) >= 32 and max(guess) <= 126:
			good.append(k)
	return good

def generate_all_keys(poss):
	all_comb = []

	count = [0 for _ in poss]
	while True:
		# Check if the count is valid
		for i in range(len(count)-1, 0, -1):
			if count[i] >= len(poss[i]):
				count[i] = 0
				count[i-1] += 1

		# If the first one is out of range return
		if count[0] >= len(poss[0]):
			return all_comb

		# Generate a key
		k = bytes([poss[i][count[i]] for i in range(len(poss))])
		all_comb.append(k)

		# Increment the last one
		count[-1] += 1

def test_key(data, k):
	plain = xor(data, k)
	if not b'sv{' in plain:
		return False

	idx = plain.index(b'sv{')

	if not b'}' in plain[idx:]:
		return False

	end_idx = idx + plain[idx:].index(b'}')
	return plain[idx:end_idx+1]

data = open('out.xor', 'rb').read()

for keylen in range(1, 10):
	print(f'{keylen = }')

	chunks = bytes_to_chunks(data, keylen)
	good = [good_keys(c) for c in chunks]

	all_keys = generate_all_keys(good)
	print(f'Found {len(all_keys)} possible keys')

	for k in all_keys:
		is_ok = test_key(data, k)
		if is_ok:
			print('----------------------')
			print(f'{k = } --> guess = {is_ok}')