import sys
import os
import math
import uuid
import struct
import configparser
from struct import unpack, pack
from binascii import crc32

if len(sys.argv) > 4 or len(sys.argv) < 2:
	print("Usage: python3 resize_user.py <path to emummc.bin file> <size to make the USER partition in GiB (float)> [path to prod.keys or bis_key_03 in hex]")
	exit()

try:
	from mbedtls import cipher
except ModuleNotFoundError:
	print("This requires python-mbedtls. You can install it via 'pip install --user python-mbedtls'")
	exit()

try:
	f = open(sys.argv[1], "r+b")
except FileNotFoundError:
	print("Please enter an valid filename")
	exit()

try:
	user_sector_size = int(float(sys.argv[2]) * (1024 ** 3) // 512)
except ValueError:
	print("Please enter a floating point number in GiB")
	exit()
	
if user_sector_size < 65536:
	print("USER must be greater than 32MiB (0.3125 GiB)")
	exit()


try:
	with open(os.path.join(os.path.expanduser('~'), '.switch', 'prod.keys'), 'r') as k:
		c = configparser.ConfigParser()
		c.read_string("[keys]\n" + k.read())
		key = bytes.fromhex(c["keys"]["bis_key_03"])
except FileNotFoundError:
	if len(sys.argv) > 3:
		try:
			with open(sys.argv[3], 'r') as k:
				c = configparser.ConfigParser()
				c.read_string("[keys]\n" + k.read())
				key = bytes.fromhex(c["keys"]["bis_key_03"])
		except FileNotFoundError:
			if len(sys.argv[3]) == 64:
				try:
					key = bytes.fromhex(sys.argv[3])
				except ValueError:
					print("Please either place a prod.keys with bis_key_03 in ~/.switch/prod.keys, or enter a valid third argument with either the location of a prod.keys or your bis_key_03")
					exit()
			else:
				print("Please either place a prod.keys with bis_key_03 in ~/.switch/prod.keys, or enter a valid third argument with either the location of a prod.keys or your bis_key_03")
				exit()
	else:
		print("Please either place a prod.keys with bis_key_03 in ~/.switch/prod.keys, or enter a valid third argument with either the location of a prod.keys or your bis_key_03")
		exit()

			


def parse_gpt_header(header):
	f.seek(0x800000 + (2 * 512) + (10 * 128))
	usr = unpack('< 16s 16s Q Q Q 72s', header)
	name = usr[5].decode('utf-16le')
	
	print("Partition type GUID:\t", uuid.UUID(bytes_le=usr[0]))
	print("Unique Partition GUID:\t", uuid.UUID(bytes_le=usr[1]))
	print("Starting LBA:\t\t", usr[2])
	print("Ending LBA:\t\t",  usr[3])
	print("Attributes:\t\t", usr[4])
	print("Name:\t\t\t", usr[5].decode('utf-16le'))
	print("Size:\t\t\t", (usr[3] - usr[2] + 1) * 512 / 1024 ** 3, "GiB")
	

def fix_user_partition_entry(size):
	# Read out gpt entry for sanity
	# parse_gpt_header(f.read(128))
	f.seek(0x800000 + 512)
	magic, = unpack("<8s", f.read(8))
	if magic != b"EFI PART":
		print("Error reading GPT header. Possible invalid emuMMC image: ", magic)
		exit()

	# calculate new end sector
	f.seek(0x800000 + (2 * 512) + (10 * 128) + 32)
	lba_start, = unpack("<Q", f.read(8))
	new_last_sector = lba_start + size - 1

	# write back new end sector
	f.seek(0x800000 + (2 * 512) + (10 * 128) + 40)
	f.write(pack('<Q', new_last_sector))
	
	return new_last_sector
	

def fix_gpt_main_header(last_sector):
### update primary header ###
	# set last usable lba
	f.seek(0x800000 + 512 + 48)
	f.write(pack('<Q', last_sector))
	# set new backup header lba
	f.seek(0x800000 + 512 + 32)
	f.write(pack('<Q', last_sector + 33))

	# get number of partitions
	f.seek(0x800000 + 512 + 80)
	num_partitions = unpack('<I', f.read(4))[0]
	# set table-crc32
	f.seek(0x800000 + 1024)
	table = f.read(128 * num_partitions)
	f.seek(0x800000 + 512 + 88)
	f.write(pack('<I', crc32(table)))

	# write primary header header-crc32
	f.seek(0x800000 + 512)
	header = bytearray(f.read(92))
	header[16] = 0
	header[17] = 0
	header[18] = 0
	header[19] = 0
	f.seek(0x800000 + 512 + 16)
	f.write(pack('<I', crc32(header)))

### make backup header ###
	# save copy of primary header
	f.seek(0x800000 + 512)
	backup_header = f.read(92)
	# write backup table
	f.seek(0x800000 + (512 * (last_sector + 30)))
	f.write(table)
	# write backup header
	f.seek(0x800000 + (512 * (last_sector + 33)))
	f.write(backup_header)
	f.write(b'0' * 420)

	# set backup header current lba
	f.seek(0x800000 + (512 * (last_sector + 33)) + 24)
	f.write(pack('<Q', last_sector + 33))
	# set backup header backup lba
	f.seek(0x800000 + (512 * (last_sector + 33)) + 32)
	f.write(pack('<Q', 1))
	# set backup header table lba
	f.seek(0x800000 + (512 * (last_sector + 33)) + 72)
	f.write(pack('<I', last_sector + 30)) # TODO: Make this not hard coded offset?
	
	# write backup header header-crc32
	f.seek(0x800000 + (512 * (last_sector + 33)))
	header = bytearray(f.read(92))
	header[16] = 0
	header[17] = 0
	header[18] = 0
	header[19] = 0
	f.seek(0x800000 + (512 * (last_sector + 33)) + 16)
	f.write(pack('<I', crc32(header)))


def parse_fat32_header(data):
	fat32 = unpack("<3x8sHBHBHHcHHHIIIHHIHH12sBBBI11s8s", data[:0x5a])
	sections = ["oem", "bytes per sector", "sectors per cluster", "number of reserved sectors", "number of fats", "max number of fat12 root dirs", "total logical sectors", "media descriptor", "logical sectors per fat for fat12", "physical sectors per track", "heads per disk", "num hidden sectors preceding FAT volume", "total logical sectors", "sectors per fat", "drive desc flags", "version", "cluster number of root dir start", "sector num of FS info", "sector num of first fat32 boot sector copy", "reserved", "physical drive number", "idk", "extended boot sig", "volume id", "volume label", "file system type"]
	for entry in zip(sections, fat32):
		print(entry)
	return data


def fix_user_fat32(size):
	fat_size = -(-(size + 32) // (4098 * 32)) * 32

	# get user partition start lba
	f.seek(0x800000 + (2 * 512) + (10 * 128) + 32)
	start_lba = unpack('<Q', f.read(8))[0]

	# decrypt fat32 header
	c = cipher.AES.new(key, cipher.MODE_XTS, pack('<QQ', 0, 0))
	f.seek(0x800000 + (start_lba * 512))
	header = bytearray(c.decrypt(f.read(0x4000)))
	# parse_fat32_header(header)
	magic, = unpack('<5s', header[0x52:0x57])
	if magic != b"FAT32":
		print("Invalid FAT32 header. Either corrupt image or incorrect bis key: ", magic)
		exit()

	# change fat32 header
	# ensure 512 bytes/sector, 32 sectors/cluster, 32 reserved sectors, 2 FAT tables
	header[0x0B:0x13] = pack('<HBHBH', 512, 32, 32, 2, 0)
	header[0x16:0x18] = pack('<H', 0)
	header[0x20:0x28] = pack('<II', size, fat_size)	# fix total logic sectors, sectors per fat
	header[0x2C:0x32] = pack('<IH', 2, 1)		# ensure proper dir cluster and fs info sector
	# fix fs info sector
	header[0x200:0x204] = pack('<4s', 	b"RRaA")
	header[0x204:0x3E4] = pack('<480B', 	*([0x00] * 480))
	header[0x3E4:0x3E8] = pack('<4s', 	b"rrAa")
	header[0x3E8:0x3F0] = pack('<8B', 	*([0xFF] * 8))
	header[0x3F0:0x3FC] = pack('<12B', 	*([0x00] * 12))
	header[0x3FC:0x400] = pack('<4B', 0x00, 0x00, 0x55, 0xAA)

	# encrypt and write fat32 header
	f.seek(0x800000 + (start_lba * 512))
	f.write(c.encrypt(header))

	# remake file allocation tables
	for i in range(3):	# 2 file allocation tables, and one more time to clear root dir
		for j in range(int(fat_size / 32)): # loop over every FAT cluster (should be 32 aligned already)
			cluster = int(1 + (i * fat_size / 32) + j)
	
			data = [0x00] * 0x4000
			if j == 0 and i < 2:
				data[0x0:0xC] = [0xf8, 0xff, 0xff, 0x0f, 
						 0xff, 0xff, 0xff, 0x0f, 
						 0xf8, 0xff, 0xff, 0x0f]
			f.seek(0x800000 + (start_lba * 512) + (cluster * 0x4000))
			c = cipher.AES.new(key, cipher.MODE_XTS, pack('>QQ', 0, cluster))
			f.write(c.encrypt(bytes(data)))
			if i == 2:	
				break
	
	
last_sector = fix_user_partition_entry(user_sector_size)
fix_gpt_main_header(last_sector)
fix_user_fat32(user_sector_size)
f.truncate(0x800000 + (last_sector + 34) * 512)		# truncate file

f.close()