from Crypto.Cipher import AES
from Crypto.Util import Counter
import struct

import hashlib
from binascii import hexlify, unhexlify

"""
typedef struct boot_dat_hdr
{
	unsigned char ident[0x10];
	unsigned char sha2_s2[0x20];
	unsigned int s2_dst;
	unsigned int s2_size;
	unsigned int s2_enc;
	unsigned char pad[0x10];
	unsigned int s3_size;
	unsigned char pad2[0x90];
	unsigned char sha2_hdr[0x20];
} boot_dat_hdr_t;
"""

def aes_ctr_dec(buf, key, iv):
    ctr = Counter.new(128, initial_value=int(hexlify(iv), 16))
    return AES.new(key, AES.MODE_CTR, counter=ctr).encrypt(buf)

boot = open('boot_recompiled.dat', 'wb')
stage2 = open('stage2_40020000 (edited with hash of databin).bin', 'rb').read()

e0sHashBytes = b""

# write ident
boot.write(b'\x54\x58\x20\x42\x4F\x4F\x54\x00\x00\x00\x00\x00\x56\x31\x2E\x30')
e0sHashBytes += b'\x54\x58\x20\x42\x4F\x4F\x54\x00\x00\x00\x00\x00\x56\x31\x2E\x30'

# write sha2-256 of stage2_40020000.bin
sha256 = hashlib.new('sha256')
sha256.update(stage2)
boot.write(sha256.digest())
e0sHashBytes += sha256.digest()

# todo: write s2_dst, hardcoded :\
boot.write(b'\x00\x00\x02\x40')
e0sHashBytes += b'\x00\x00\x02\x40'

# write s2_size
boot.write(struct.pack('I', len(stage2)))
e0sHashBytes += struct.pack('I', len(stage2))

# write s2_enc
boot.write(struct.pack('I', 1))
e0sHashBytes += struct.pack('I', 1)

# 0x10 size padding
boot.write(b'\x00' * 0x10)
e0sHashBytes += b'\x00' * 0x10

# s3_size?
boot.write(b'\x50\x2B\xED\x00')
e0sHashBytes += b'\x50\x2B\xED\x00'

# 0x90 size padding
boot.write(b'\x00' * 0x90)
e0sHashBytes += b'\x00' * 0x90

# calculate e0ssha256
sha256 = hashlib.new('sha256')
sha256.update(e0sHashBytes)
boot.write(sha256.digest())

# stage2 section
boot.write(aes_ctr_dec(stage2, unhexlify("47E6BFB05965ABCD00E2EE4DDF540261"), unhexlify("8E4C7889CBAE4A3D64797DDA84BDB086")))

# data section
with open('data_80000000 (edited with pub key).bin', 'rb') as fh:
	boot.write(aes_ctr_dec(fh.read(), unhexlify("030D865B7E458B10AD5706F6E227F4EB"), unhexlify("AFFC93692EBD2E3D252339F01E03416B")))

# fb section
with open('fb_F0000000.bin', 'rb') as fh:
	boot.write(aes_ctr_dec(fh.read(), unhexlify("E2AC05206A701C9AA514D2B2B7C9F395"), unhexlify("46FAB59AF0E469EF116614DEC366D15F")))

# write arm64
with open('arm64_80FFFE00.bin', 'rb') as fh:
	boot.write(aes_ctr_dec(fh.read(), unhexlify("35D8FFC4AA1BAB9514825EB0658FB493"), unhexlify("C38EA26FF3CCE98FD8D5ED431D9D5B94")))

# write rest of boot.dat og from 0x571e20 onwards
with open('boot.dat', 'rb') as fh:
	fh.seek(0x571E20, 0)
	boot.write(fh.read())

boot.close()