# Modified path ways by Swizzy to fit xeBuild GUI better :)
# Made the script read in custom keys (1blkey/cpukey) from xeBuild GUI
# Original script (before my mods) was made by TX RGH Team

# modified by GliGli and Tiros for the reset glitch hack

# you might need to fill in this
import os.path
secret_1BL = "\xDD\x88\xAD\x0C\x9E\xD6\x69\xE7\xB5\x67\x94\xFB\x68\x56\x3E\xFA"
cpukey = None

# cpukey is needed when supplying a non-14717 nand image
#cpukey = 


Xell_Path = "xell-gggggg.bin"
Xell_CD_Path = "CD\\CD2.0"
CB_Path = "CB\\" # Added this for ease of changing path for CB files laters...

XELL_BASE_FLASH = 0xc0000
CODE_BASE = 0x1c000000

SCRIPT_VERSION = 0x01

Keyvault = None
SMC = None
CB_A = None
CB_A_img = None

CB_A_crypted = None
CB_B = None
CD = None
CD_plain = None
CE = None
CF = None
CG = None
Xell = ""
Exploit = None

#if secret_1BL is None:
#	secret_1BL = open("key_1BL.bin", "rb").read()

if os.path.isfile("1blkey.bin"):
        secret_1BL = open("1blkey.bin", "rb").read()
if os.path.isfile("cpukey.bin"):
        cpukey = open("cpukey.bin", "rb").read()

 # Import Psyco if available
try:
	import psyco
	psyco.full()
except ImportError:
	pass

# first, unpack base input image. We are ignoring any updates here
import hmac, sha, struct, sys, binascii
try:
	import Crypto.Cipher.ARC4 as RC4
except ImportError:
	print "Error importing Crypto.Cipher.ARC4 - please install python-crypto!"
	print "You can get it from http://www.dlitz.net/software/pycrypto/"
	sys.exit(-1)

def unpack_base_image(image):
	global SMC, CB_A, CB_B, CD, CE, Keyvault

	if image[0x205] == "\xFF" or image[0x415] == "\xFF" or image[0x200] == "\xFF":
		print "ECC'ed - will unecc."
		res = ""
		for s in range(0, len(image), 528):
			res += image[s:s+512]
		image = res

	unpackstring = "!HHLLL64s5LLLLLLLL"
	(id1, build, flags, bloffset, size0, copyright, z0, z1, z2, z3, r7, size1, r3, r4, z5, z6, smc_len, smc_start) = struct.unpack(unpackstring, image[:struct.calcsize(unpackstring)])
	#assert not (z0 or z1 or z2 or z3 or z5 or z6), "zeros are not zero."

	block_offset = bloffset

	SMC = image[smc_start:smc_start+smc_len]
	Keyvault = image[0x4000:0x8000]

	#assert smc_len == 0x3000, "never saw an SMC != 0x3000 bytes"
	semi = 0
	for block in range(30):
		(block_id, block_build, block_flags, block_entry_point, block_size) = struct.unpack("!2sHLLL", image[block_offset:block_offset+16])
		block_size += 0xF
		block_size &= ~0xF
		id = ord(block_id[1]) & 0xF

		print "Found %dBL (build %d) at %08x" % (id, block_build, block_offset)
		data = image[block_offset:block_offset+block_size]
		
		if id == 2:
			if semi == 0:
				CB_A = data
				semi = 1
			elif semi ==1:
				CB_B = data
				semi = 0

		elif id == 4:
			CD = data
		elif id == 5:
			CE = data
		
		block_offset += block_size

		if id == 5:
			break

	assert CB_A and CD

def unpack_update(image):
	global CF, CG
	
	block_offset = 0
	for block in range(30):
		(block_id, block_build, block_flags, block_entry_point, block_size) = struct.unpack("!2sHLLL", image[block_offset:block_offset+16])
		block_size += 0xF
		block_size &= ~0xF
		id = ord(block_id[1]) & 0xF

		print "Found %dBL (build %d) at %08x" % (id, block_build, block_offset)
		data = image[block_offset:block_offset+block_size]
		
		if id == 6:
			CF = data
		elif id == 7:
			CG = data

		block_offset += block_size

		if id == 7:
			break

def build(data):
	return struct.unpack(">H", data[2:4])[0]

def decrypt_CB(CB):
	secret = secret_1BL
	key = hmac.new(secret, CB[0x10:0x20], sha).digest()[0:0x10]
	CB = CB[0:0x10] + key + RC4.new(key).decrypt(CB[0x20:])
	return CB

def decrypt_CB_Cpu(CB):
	assert cpukey
	secret = CB_A[0x10:0x20]
	h = hmac.new(secret,None, sha);
	h.update(CB[0x10:0x20]);
	h.update(cpukey);
	key = h.digest()[0:0x10]
	CB = CB[0:0x10] +key+ RC4.new(key).decrypt(CB[0x20:])
	return CB

def decrypt_CD(CD, CB):
# enable this code if you want to extract CD from a flash image and you know the cup key.
# disable this when this is a zero-paired image.
#	assert cpukey or build(CD) < 1920
	secret = CB[0x10:0x20]
	key = hmac.new(secret, CD[0x10:0x20], sha).digest()[0:0x10]
#	if build(CD) >= 1920:
#		key = hmac.new(cpukey, key, sha).digest()[0:0x10]
	CD = CD[0:0x10] + key + RC4.new(key).decrypt(CD[0x20:])
	return CD

def decrypt_CE(CE, CD):
	secret = CD[0x10:0x20]
	key = hmac.new(secret, CE[0x10:0x20], sha).digest()[0:0x10]
	CE = CE[0:0x10] + key + RC4.new(key).decrypt(CE[0x20:])
	return CE

def decrypt_CF(CF):
	secret = secret_1BL
	key = hmac.new(secret, CF[0x20:0x30], sha).digest()[0:0x10]
	CF = CF[0:0x20] + key + RC4.new(key).decrypt(CF[0x30:])
	return CF

def decrypt_CG(CG, CF):
	secret = CF[0x330:0x330+0x10]
	key = hmac.new(secret, CG[0x10:0x20], sha).digest()[0:0x10]
	CG = CG[:0x10] + key + RC4.new(key).decrypt(CG[0x20:])
	return CG

def decrypt_SMC(SMC):
	key = [0x42, 0x75, 0x4e, 0x79]
	res = ""
	for i in range(len(SMC)):
		j = ord(SMC[i])
		mod = j * 0xFB
		res += chr(j ^ (key[i&3] & 0xFF))
		key[(i+1)&3] += mod
		key[(i+2)&3] += mod >> 8
	return res

def encrypt_CB(CB, random):
	secret = secret_1BL
	key = hmac.new(secret, random, sha).digest()[0:0x10]
	CB = CB[0:0x10] + random + RC4.new(key).encrypt(CB[0x20:])
	return CB, key

def encrypt_CB_Cpu(CB_B, CB_A_key):
	assert cpukey
	secret = CB_A_key
	crypto = CB_B[0x10:0x20]

	key = hmac.new(secret, crypto + cpukey, sha).digest()[0:0x10]
	
	CB_B = CB_B[0:0x10] + crypto + RC4.new(key).encrypt(CB_B[0x20:])
	return CB_B, key

def encrypt_CD(CD, CB_key, random):
	secret = CB_key
	key = hmac.new(secret, random, sha).digest()[0:0x10]
	CD = CD[0:0x10] + random + RC4.new(key).encrypt(CD[0x20:])
	return CD, key

def encrypt_CE(CE, CD_key, random):
	secret = CD_key
	key = hmac.new(secret, random, sha).digest()[0:0x10]
	CE = CE[0:0x10] + random + RC4.new(key).encrypt(CE[0x20:])
	return CE

def encrypt_CF(CF, random):
	secret = secret_1BL
	key = hmac.new(secret, random, sha).digest()[0:0x10]
	CF_key = CF[0x330:0x330+0x10]
	CF = CF[0:0x20] + random + RC4.new(key).encrypt(CF[0x30:])
	return CF, CF_key

def encrypt_CG(CG, CF_key, random):
	secret = CF_key
	key = hmac.new(secret, random, sha).digest()[0:0x10]
	CG = CG[:0x10] + random + RC4.new(key).encrypt(CG[0x20:])
	return CG

def encrypt_SMC(SMC):
	key = [0x42, 0x75, 0x4e, 0x79]
	res = ""
	for i in range(len(SMC)):
		j = ord(SMC[i]) ^ (key[i&3] & 0xFF)
		mod = j * 0xFB
		res += chr(j)
		key[(i+1)&3] += mod
		key[(i+2)&3] += mod >> 8
	return res

# CB_patches is an array of patchsets, a patchset is a version number followed by an array of patches, a patch is 3 ints: [offset,plaintext,patch]
CB_patches = [[9188,[[0x4f08,0x409a0010,0x60000000],[0x5618,0x480018e1,0x60000000],[0x5678,0x480000b9,0x60000000], [0x4d10, 0x7BEB0620, 0x48000168]]],
[13121,[[0x5048,0x7BEB0620,0x48000168],[0x5240,0x409A0010,0x60000000],[0x58ac,0x480001C5,0x60000000],[0x5958,0x480018E1,0x60000000],[0x59B8,0x480000B9,0x60000000]]],
[6752,[[0x6AA0,0x409A0010,0x60000000],[0x71B0,0x480018D9,0x60000000],[0x7200,0x419A0014,0x48000014],[0x68A8,0x7BEB0620,0x48000168]]],
[5772,[[0x6A58,0x409A0010,0x60000000],[0x7168,0x480018E1,0x60000000],[0x71b8,0x419A0014,0x48000014],[0x6860,0x7BEB0620,0x48000168],[0x4F0,0x48006239,0x60000000]]],
[4577,[[0x53C0,0x7BEB0620,0x48000168],[0x55B8,0x409A0010,0x60000000],[0x5D54,0x480018C5,0x60000000],[0x5DB0,0x480000A1,0x60000000]]]]

def int_to_str(i):
	return [chr((i>>24) & 0xff),chr((i>>16) & 0xff),chr((i>>8) & 0xff),chr(i & 0xff)]

def patch_CB(CB):
	found = False

	for versions in CB_patches:
		if build(CB) == versions[0]:
			print "patchset for %d found, %d patch(es)" % (versions[0],len(versions[1]))
			found  = True
			for patches in versions[1]:
				plain = int_to_str(patches[1])
				patch = int_to_str(patches[2])

				patched = ""

				for i in range(4):
					keystream = ord(plain[i]) ^ ord(CB[i+patches[0]])
					patched = patched + chr(keystream ^ ord(patch[i]))

				CB = CB[:patches[0]] + patched + CB[patches[0]+4:]

	#assert found,"can't patch that CB"

	return CB
	


def xor_hack(CB_B, CB_B_plain, CB_B_patched):
	assert CB_B, "You need a CB_B!"
	assert CB_B_plain, "You need a decrypted CB_B for build %d" % (build(CB_B))
	assert CB_B_patched, "You need a patched exploitable CB_B!"
	
	headerlen = 0x40
	offset = headerlen
	
	CBKeystream = ""
	
	assert (len(CB_B_patched) <= len(CB_B)), "This CB (%X) is too large to fit inside your CBB (%X)!" % (len(CB_B_patched), len(CB_B))
		
	for ctr in range((len(CB_B_patched) - headerlen)/4):
		
		plain = struct.unpack(">L", CB_B_plain[offset:offset+4])[0]
		patch = struct.unpack(">L", CB_B_patched[offset:offset+4])[0]
		
		plain = int_to_str(plain)
		patch = int_to_str(patch)
		
		patched = ""

		for i in range(4):
			keystream = ord(plain[i]) ^ ord(CB_B[i+offset])
			patched = patched + chr(keystream ^ ord(patch[i]))
			CBKeystream = CBKeystream + chr(keystream)

		CB_B = CB_B[:offset] + patched + CB_B[offset+4:]
		
		offset = offset+4
	
	#open("output/keystream.bin", "wb").write(CBKeystream)
	
	CB_B= CB_B_patched[:0xC] + CB_B[0xC:]

	return CB_B

	
#Dynamic SMC Patching by cOz
console_types = ["none/unk","Xenon","Zephyr","Falcon","Jasper","Trinity","Corona","Winchester"]

def patch_SMC(SMC):
	found = False
	smctyp = (ord(SMC[0x100])>>4)&0xF;
	for bytes in range((len(SMC)-8)):
		if ord(SMC[bytes]) == 0x05:
			if ((ord(SMC[bytes+2]))==0xE5) and ((ord(SMC[bytes+4]))==0xb4) and ((ord(SMC[bytes+5]))==0x05):
				found = True
				print "Patching %s version %d.%d SMC at offset 0x%x" % (console_types[smctyp], ord(SMC[0x101]), ord(SMC[0x102]), bytes)
				SMC = SMC[:bytes] + chr(0x0) + chr(0x0) + SMC[bytes+2:]
	if not found:
		print " ! Warning: can't patch this %s type SMC!" % (console_types[smctyp])
	return SMC

	
def consoleType(SMC):
	
	smctyp = (ord(SMC[0x100])>>4)&0xF;
	
	return console_types[smctyp]

def allzero(string):
	for x in string:
		if ord(x):
			return False
	return True

def allFF(string):
	for x in string:
		if ord(x) != 0xFF:
			return False
	return True

def calcecc(data):
	assert len(data) == 0x210
	val = 0
	for i in range(0x1066):
		if not i & 31:
			v = ~struct.unpack("<L", data[i/8:i/8+4])[0]
		val ^= v & 1
		v >>= 1
		if val & 1:
			val ^= 0x6954559
		val >>= 1

	val = ~val
	return data[:-4] + struct.pack("<L", (val << 6) & 0xFFFFFFFF)

def addecc(data, block = 0, off_8 = "\x00" * 4):
	res = ""
	while len(data):
		d = (data[:0x200] + "\x00" * 0x200)[:0x200]
		data = data[0x200:]
		
		d += struct.pack("<L4B4s4s", block / 32, 0, 0xFF, 0, 0, off_8, "\0\0\0\0")
		d = calcecc(d)
		block += 1
		res += d
	return res

import sys

print "\nUnpatchable Universal XELL Nand Image Builder\n"

for i in sys.argv[1:]:
	image = open(i, "rb").read()[:1*1024*1024]
	if image[:2] == "\xFF\x4F":
		print " * found flash image, unpacking..."
		unpack_base_image(image)
		CB_A_crypted = CB_A
		
		#print " * not decrypting SMC"
		SMC = decrypt_SMC(SMC)
	elif image[:2] == "CD" and allzero(image[0x20:0x230]):
		print " * found decrypted CD"
		CD_plain = image
	elif len(image) == 0x3000 and image.find("<Copyright 2001-") >= 0:
		print " * found decrypted SMC"
		SMC = image
	elif len(image) == 0x3000:
		print " * found encrypted SMC (i hope so)"
		SMC = decrypt_SMC(image)
	elif image[-0x10:] == "x" * 16:
		print " * found XeLL binary, must be linked to %08x" % CODE_BASE

		assert len(image) <= 256*1024
		image = (image + "\0" * 256*1024)[:256*1024]
		Xell += image
	else:
		raise " * unknown image found in file %s!" % i
		
#xenon_builds = []
zephyr_builds = [4578, 4575, 4577]
falcon_builds = [5771, 5772, 5773]
jasper_builds = [6750, 6752, 6753]
trinity_builds = [9188, 9230]
corona_builds = [13121, 13180]



xor_hack_builds = [9230, 5773, 6753, 4575, 13180]

patch_builds = [5772, 6572, 9188, 13121]


print "\n * we found the following parts:"
print "SMC: %d.%d" %(ord(SMC[0x101]),ord(SMC[0x102]))
print "CB_A:", CB_A and build(CB_A) or "missing"
print "CB_B:", CB_B and build(CB_B) or "missing"
print "CD (image):", CD and build(CD) or "missing"
print "CD (decrypted):", CD_plain and build(CD_plain) or "missing"
print ""


print " * checking for proper 1BL key...",
sum = 0
for v in secret_1BL:
	sum = sum + ord(v)
assert sum == 0x983, "you need to fill in secret_1BL properly!"
print "ok"

CB_A_img_RAND = CB_A_crypted[0x10:0x20]
conType = consoleType(SMC)

print " * decrypting..."
CB_A_img = decrypt_CB(CB_A_crypted)

print " * checking if all files decrypted properly...",
#assert allzero(SMC[-4:])
print "ok"

print " * checking required versions..."
assert build(CB_A) in falcon_builds or build(CB_A) in zephyr_builds or build(CB_A) in jasper_builds or build(CB_A) in trinity_builds or build(CB_A) in corona_builds, "CB "+str(build(CB_A)) + " not eligible"
assert build(CB_A) in xor_hack_builds or patch_builds or cpukey, "Image must have a split CB or you must supply cpu_key!"
if not CD_plain:
	print "  ** Decrypted CD not found! Loading xell CD from %s" % (Xell_CD_Path)
	CD_plain = open(Xell_CD_Path, "rb").read()
	
if not Xell:
	print "  ** Xell binary not found! Loading xell from %s" % (Xell_Path)
	
	image = open(Xell_Path, "rb").read()[:1*1024*1024]
	assert len(image) <= 256*1024 
	assert image[-0x10:] == "x" * 16
	
	print "    * found XeLL binary, must be linked to %08x" % CODE_BASE
	image = (image + "\0" * 256*1024)[:256*1024]
	Xell += image
	
print "ok\n"

Final = ""

print " * patching SMC..."
SMC=patch_SMC(SMC)
#print " * not patching SMC"

print " * Replacing CD..."
CD = CD_plain

cbapath = ""
cbbpath = ""
	
if build(CB_A) in falcon_builds: 
	cbapath = CB_Path + "cba_5772.bin"
	cbbpath = CB_Path + "cbb_5772.bin"
elif build(CB_A) in jasper_builds: 
	cbapath = CB_Path + "cba_6752.bin"
	cbbpath = CB_Path + "cbb_6752.bin"
elif build(CB_A) in trinity_builds: 
	cbapath = CB_Path + "cba_9188.bin"
	cbbpath = CB_Path + "cbb_9188.bin"
	CD = open("CD\\CD", "rb").read()
elif build(CB_A) in zephyr_builds: 
	cbapath = CB_Path + "cba_4577.bin"
	cbbpath = CB_Path + "cbb_4577.bin"
elif build(CB_A) in corona_builds: 
	cbapath = CB_Path + "cba_13121.bin"
	cbbpath = CB_Path + "cbb_13121.bin"
	CD = open("CD\\CD", "rb").read()
	
if not cpukey: 

	if build(CB_B) in patch_builds:
		print " * patching CB_B " + str(build(CB_B)) + "..."
		CB_B = patch_CB(CB_B)
		CB_A = CB_A_crypted
		
	if build(CB_B) in xor_hack_builds:
		print "\n * XOR HACK NEEDED FOR %s CB %d" % (conType, build(CB_B))
		
		cbbplainpath = CB_Path + "cbb_" + str(build(CB_B)) + ".bin"
		
		CB_B_plain = open(cbbplainpath, "rb").read()
		CB_B_patched = open(cbbpath, "rb").read()
		
		print " * patching CB_B " + str(build(CB_B_patched)) + "..."
		CB_B_patched = patch_CB(CB_B_patched)
			
		print " * Applying XOR Hack to CB_B %d" % (build(CB_B))
		CB_B = xor_hack(CB_B, CB_B_plain, CB_B_patched)
	
		print " * Replacing CB_A %d with %d" % (build(CB_A), build(CB_B_patched))
		CB_A = open(cbapath, "rb").read()
		
		CB_A, CB_A_key = encrypt_CB(CB_A, CB_A_img_RAND)
			
		print ""
		
else: 
	print "\n * Building new %s bootloader chain using cpu_key: " % (conType)
	hexdata = ''.join('%02X' % ord(byte) for byte in cpukey)
	print " " + hexdata
	
	CB_A = open(cbapath, "rb").read()
	CB_B = open(cbbpath, "rb").read()
	
	CB_B = patch_CB(CB_B)
	
	CB_A, CB_A_key = encrypt_CB(CB_A, CB_A_img_RAND)
	CB_B, CB_A_key = encrypt_CB_Cpu(CB_B, CB_A_key)
	
	
c = "RGH2 2stage CB img"

	
print " * constructing new image..."

base_size = 0x8000 + len(CB_A) + len(CD) + len(CE)
if CB_B:
	base_size += len(CB_B)
base_size += 16383
base_size &=~16383
patch_offset = base_size

SMC_size = len(SMC)
SMC_offset = 0x4000-SMC_size

print " * base size: %x" % base_size

c += ", version=%02x, CB=%d" % (SCRIPT_VERSION, build(CB_A))
Header = struct.pack(">HHLLL64s5LLLLLLLL", 0xFF4F, 1888, 0, 0x8000, base_size, c, 0, 0, 0, 0, 0x4000, patch_offset, 0x20712, 0x4000, 0, 0, SMC_size, SMC_offset)

Header = (Header + "\0" * 0x1000)[:0x1000]
random = "\0" * 16

#print " * not re-encrypting SMC (already encrypted)"
SMC = encrypt_SMC(SMC)

	
assert len(Final) < XELL_BASE_FLASH, "Please move XELL_BASE_FLASH"

if len(Xell) <= 256*1024:
	print " * No separate recovery Xell available!"
	Xell *= 2

print "\n * %s Image Flash Layout:" % (conType)

def add_to_flash(d, w, offset = 0):
	global Final
	print "0x%08x..0x%08x (0x%08x bytes) %s" % (offset + len(Final), offset + len(Final) + len(d) - 1, len(d), w)
	Final += d

def pad_to(loc):
	global Final
	pad = "\xFF" * (loc - len(Final))
	add_to_flash(pad, "Padding")

add_to_flash(Header[:0x200], "Header")
pad_to(SMC_offset)
add_to_flash(SMC, "SMC")
add_to_flash(Keyvault, "Keyvault")


add_to_flash(CB_A, "CB_A %d" % build(CB_A))
add_to_flash(CB_B, "CB_B %d" % build(CB_B))

add_to_flash(CD, "CD %d" % build(CD))

pad_to(XELL_BASE_FLASH)
add_to_flash(Xell[0:256*1024], "Xell (backup)")
add_to_flash(Xell[256*1024:], "Xell (main)")

print " * Encoding ECC..."
Final = addecc(Final)

open("output/image_00000000.ecc", "wb").write(Final)
print "------------- Written into output/image_00000000.ecc"
