import os,sys,binascii,hashlib

# known plaintexts for, repectively, native firm3.0+, b9s 1.0-1.3, b9s 1.4, fb3ds 1.1+, luma 8.0-11.0
sectors = (2, 11, 1, 4, 5)

s2firm_name = "SafeB9SInstaller"
s2firm_max_size = 0x0B930000 - 0x0B800000
payload_name = "boot9strap"

def generateHeader(buff, arrayName, ifsha256=False, sectors=None):
	f=open(sourcedir+"/"+arrayName+".h","w")
	f.write("#pragma once\n\n");
	f.write("const size_t %s_len = %d;\n" % (arrayName, len(buff)))
	f.write("const u8 "+arrayName+"_hash[] = {"+buff2array(hash(buff, ifsha256))+"};\n")
	f.write("static u8 "+arrayName+"[] __attribute__((aligned(4))) = {\n")
	f.write(buff2array(buff))
	f.write("\n")
	f.write("};\n")
	if sectors is not None:
		f.write("const size_t sector_offsets[] = %s;\n" % str(sectors).replace('(', '{ ').replace(')', ' }'))
		f.write("const size_t num_sectors = %d;\n" % len(sectors))
	f.close()

def hash(buf,ifsha256=False):
	obj=None
	if ifsha256:
		return hashlib.sha256(buf).digest()
	return hashlib.sha1(buf).digest()

def buff2array(s):
	final=["",] * len(s)
	for i in range(len(s)):
		final[i] = "0x%02X" % s[i]
	return ",".join(final)

if len(sys.argv) != 2:
	print("Usage: build_headers.py <build phase>")
	sys.exit(1)
if sys.argv[1] != "1" and sys.argv[1] != "2":
	print("No such phase:", sys.argv[1])
	sys.exit(1)

phase2 = (sys.argv[1] == "2")
sourcedir = "source/bins"
if not phase2:
	sourcedir = "stage2_installer/" + sourcedir

try:
	os.mkdir(sourcedir)
except:
	pass

if not phase2:
	if os.path.exists("embedded_files/secret_sector.bin"):
		buff_ssec = open("embedded_files/secret_sector.bin", 'rb').read(-1)
		if hash(buff_ssec) != b'\xb9\xcf\xa8\x49\x16\xa9\x30\xd2\x72\x25\x01\x30\xd6\x7c\xeb\x82\x21\x41\x17\x7d':
			print("Invalid secret_sector.bin detected, will not be embedded in stage2 installer")
		else:
			generateHeader(buff_ssec, "secret_sector", True)
	else:
		print("embedded_files/secret_sector.bin not found, will not be embedded in stage2 installer")

	buff_pay = open("embedded_files/"+payload_name+".firm", 'rb').read(-1)
	generateHeader(buff_pay, "payload", True)

else:
	if not (1 <= len(sectors) <= 5):
		print("Invalid number of sectors")
		sys.exit(1)

	buff_secs = b''

	for sector in sectors:
		if sector < 0:
			print("huh?")
			sys.exit(1)
		buff = open("embedded_files/%d.bin" % sector, 'rb').read(-1)
		if len(buff) != 0x200:
			print("Unexpected sector %d size: 0x%X" % (sector, len(buff)))
			sys.exit(1)
		buff_secs += buff

	if len(buff_secs) != len(sectors)*0x200:
		print("wut?")
		sys.exit(1)

	b9prot = open("embedded_files/boot9.bin", 'rb').read(-1)[-0x8000:]
	if hash(b9prot) != b'\x74\x1A\xCB\x62\x9A\xBB\x67\xF6\xD9\xB7\x08\x18\x4F\xBD\x9E\x01\xEC\x2F\xCE\x3C':
		print("Invalid boot9.bin provided")
		sys.exit(1)

	buff_boot9_blob =  b9prot[0x56E0:0x5700] # OTP normal key + IV
	buff_boot9_blob += b9prot[0x5A50:0x5A60] # slot 4 keyY
	buff_boot9_blob += b9prot[0x5A70:0x5A80] # slot 6 keyY
	buff_boot9_blob += b9prot[0x5860:0x58A4] # data + IV + more data for keygen

	buff_s2f = open("stage2_installer/output/"+s2firm_name+".firm", 'rb').read(-1)
	if len(buff_s2f) > s2firm_max_size:
		print("Stage 2 installer too big!")
		sys.exit(1)

	buff_min = open("embedded_files/min.firm", 'rb').read(-1)
	if len(buff_min) != 0x200:
		print("Unexpected minfirm size: 0x%X" % len(buff_min))
		sys.exit(1)

	generateHeader(buff_boot9_blob, "boot9_blob")
	generateHeader(buff_secs, "plaintext_sectors", sectors=sectors)
	generateHeader(buff_s2f, "stage2_firm")
	generateHeader(buff_min, "minfirm")
