#!/usr/bin/env python3

from pwn import *
import random
import string
import math
import threading

from Crypto.Cipher import AES

#context.log_level = "debug"
max_measurements = 20000

allowed_chars = string.ascii_letters + string.digits + string.punctuation

sbox = (
	0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5, 0x30, 0x01, 0x67, 0x2B, 0xFE, 0xD7, 0xAB, 0x76,
	0xCA, 0x82, 0xC9, 0x7D, 0xFA, 0x59, 0x47, 0xF0, 0xAD, 0xD4, 0xA2, 0xAF, 0x9C, 0xA4, 0x72, 0xC0,
	0xB7, 0xFD, 0x93, 0x26, 0x36, 0x3F, 0xF7, 0xCC, 0x34, 0xA5, 0xE5, 0xF1, 0x71, 0xD8, 0x31, 0x15,
	0x04, 0xC7, 0x23, 0xC3, 0x18, 0x96, 0x05, 0x9A, 0x07, 0x12, 0x80, 0xE2, 0xEB, 0x27, 0xB2, 0x75,
	0x09, 0x83, 0x2C, 0x1A, 0x1B, 0x6E, 0x5A, 0xA0, 0x52, 0x3B, 0xD6, 0xB3, 0x29, 0xE3, 0x2F, 0x84,
	0x53, 0xD1, 0x00, 0xED, 0x20, 0xFC, 0xB1, 0x5B, 0x6A, 0xCB, 0xBE, 0x39, 0x4A, 0x4C, 0x58, 0xCF,
	0xD0, 0xEF, 0xAA, 0xFB, 0x43, 0x4D, 0x33, 0x85, 0x45, 0xF9, 0x02, 0x7F, 0x50, 0x3C, 0x9F, 0xA8,
	0x51, 0xA3, 0x40, 0x8F, 0x92, 0x9D, 0x38, 0xF5, 0xBC, 0xB6, 0xDA, 0x21, 0x10, 0xFF, 0xF3, 0xD2,
	0xCD, 0x0C, 0x13, 0xEC, 0x5F, 0x97, 0x44, 0x17, 0xC4, 0xA7, 0x7E, 0x3D, 0x64, 0x5D, 0x19, 0x73,
	0x60, 0x81, 0x4F, 0xDC, 0x22, 0x2A, 0x90, 0x88, 0x46, 0xEE, 0xB8, 0x14, 0xDE, 0x5E, 0x0B, 0xDB,
	0xE0, 0x32, 0x3A, 0x0A, 0x49, 0x06, 0x24, 0x5C, 0xC2, 0xD3, 0xAC, 0x62, 0x91, 0x95, 0xE4, 0x79,
	0xE7, 0xC8, 0x37, 0x6D, 0x8D, 0xD5, 0x4E, 0xA9, 0x6C, 0x56, 0xF4, 0xEA, 0x65, 0x7A, 0xAE, 0x08,
	0xBA, 0x78, 0x25, 0x2E, 0x1C, 0xA6, 0xB4, 0xC6, 0xE8, 0xDD, 0x74, 0x1F, 0x4B, 0xBD, 0x8B, 0x8A,
	0x70, 0x3E, 0xB5, 0x66, 0x48, 0x03, 0xF6, 0x0E, 0x61, 0x35, 0x57, 0xB9, 0x86, 0xC1, 0x1D, 0x9E,
	0xE1, 0xF8, 0x98, 0x11, 0x69, 0xD9, 0x8E, 0x94, 0x9B, 0x1E, 0x87, 0xE9, 0xCE, 0x55, 0x28, 0xDF,
	0x8C, 0xA1, 0x89, 0x0D, 0xBF, 0xE6, 0x42, 0x68, 0x41, 0x99, 0x2D, 0x0F, 0xB0, 0x54, 0xBB, 0x16
	)

def gather_measurements(remotes, amount):
	progress = log.progress("Gathering measurements")
	measurements = {}
	threads = []
	lock = threading.Lock()
	for r in remotes:
		thread = threading.Thread(target=gather_measurements_, args=(r, amount // len(remotes), amount, measurements, progress, lock))
		thread.start()
		threads.append(thread)
	
	for thread in threads:
		thread.join()
	
	progress.success("{amount}/{amount}".format(amount=amount))
	return measurements

def gather_measurements_(r, amount, total_amount, measurements, progress, lock):
	for i in range(amount):
		r.send(b"profiler_reset\n")
		r.recvuntil(b"> ")
		password = ''.join(random.choices(allowed_chars, k=16)).encode("ascii")
		r.send(b"login\n")
		r.send("admin")
		r.send(b"\n")
		r.send(password)
		r.send(b"\n")
		r.recvuntil(b"> ")
		r.send(b"profiler_print\n")
		response = r.recvuntil(b"\n> ", drop=True)
		response = response.decode("ascii")
		for line in response.splitlines():
			line = line.split(" ")
			name = line[0]
			no_invocations = int(line[-1])
			if name == "gf_reduce":
				lock.acquire()
				measurements[password] = no_invocations
				progress.status("{}/{}".format(len(measurements), total_amount))
				lock.release()
				break


def mean(data):
	return sum(data) / len(data)

def variance(data, mean):
	result = 0
	for value in data:
		result += (value - mean) ** 2
	return result / len(data)

def t_test(group_big, group_small):
	mean_big = mean(group_big)
	variance_big = variance(group_big, mean_big)
	mean_small = mean(group_small)
	variance_small = variance(group_small, mean_small)
	return (mean_big - mean_small) / math.sqrt(variance_big / len(group_big) + variance_small / len(group_small))

remotes = [process("/home/manuel/wolke/Projects/secutech_authenticator/build/meson.debug.linux.x86_64/secutech", cwd="/home/manuel/wolke/Projects/secutech_authenticator")]
#remotes = []
#no_threads = 50
#for i in range(no_threads):
#	remotes.append(remote("ccn.li", "5555"))

for r in remotes:
	r.recvuntil(b"> ")

measurements = {}

key = []

# We keep gathering measurements until we are certain enoguh which key the correct one is
while len(key) < 16:
	if len(measurements) < max_measurements:
		measurements.update(gather_measurements(remotes, 1000))
	else:
		log.failure("Unable to restore key after using the maximum configured number of measurements. This can mean one of two things: Either we just got very unlucky with the random numbers and this error will not occur again during the next run. In this case this error can be ignored. If this error keeps occuring the challange is broken.")
		sys.exit(1)
	log.info("Total number of unique measurements gathered: {}".format(len(measurements)))
	
	# This attack allows us to test each aes key byte independently
	for key_byte in range(len(key), 16):
		progress = log.progress("Attacking key byte {:2}".format(key_byte))
		
		t_values = []
		
		# We guess each possible value for the key byte. We then calculate the key addition and the byte subsitution for that byte.
		# (We leave out shift rows because shifting the rows does not change the timing)
		# Then we predict whether the calulation will be slow or fast based on that result and check which of the key guesses best fit our prediction.
		# That Guess is most likely the correct value for the key byte.
		for key_guess in range(0x100):
			group_fast = []
			group_slow = []
			for password, no_invocations in measurements.items():
				# If the msb is set the reduction function has to be called making code execution slower
				is_fast = sbox[password[key_byte] ^ key_guess] & 0x80 == 0
				# Sort the measurement into one of the groups based on our prediction
				if is_fast:
					group_fast.append(no_invocations)
				else:
					group_slow.append(no_invocations)

			# The further apart the measurements in these groups are the more likely it is that we found the correct value for the key byte.
			# We can calculate the certainty using Welch's t-test. If the result of the t-test is >4.5 we can be very certain we found the correct value.
			t_values.append(t_test(group_slow, group_fast))
		
		# The biggest value in the list is our best guess for the value
		max_t_value = max(t_values)
		candidate = t_values.index(max_t_value)
		progress.success("{:02X} ({:})".format(candidate, max_t_value))

		# Check if we are certain enough to add this result to the key. If not stop attacking and gather more measurements
		if max_t_value >= 4.5:
			key.append(candidate)
		else:
			break

encrypted_adminpw = bytes([0x95, 0x22, 0x28, 0xf3, 0x90, 0xa6, 0x0b, 0xd2, 0x5d, 0x61, 0xdd, 0x1e, 0xdf, 0x39, 0x44, 0x7b])
aes = AES.new(bytes(key), AES.MODE_ECB)
adminpw = aes.decrypt(encrypted_adminpw)
remotes[0].send("login\nadmin\n")
remotes[0].send(adminpw)
remotes[0].send("\n")
remotes[0].recvuntil("Password: ")
flag = remotes[0].recvuntil("\n")

for r in remotes:
	r.close()

print("Key:", " ".join(["{:02X}".format(k) for k in key]))
print("Adminpw:", adminpw.decode("ascii"))
print("Flag:", flag.decode("ascii"))
