from math import ceil
import os
import secrets
import struct
import zlib

from Crypto.Cipher import AES
from Crypto.Hash import HMAC, SHA1
from construct import Bytes, Const, Int16ul, Int32ul, Int64ul, Int8ul, Struct

# ---- Configuration
ENCRYPTION_KEY = bytes.fromhex("")
INPUT_FILE = ""  # Should not be encrypted.
OUTPUT_FILE = ""
BOOTID = {
    "type": 0x0201,  # 0x0201: Option, 0x0000: Pack, 0x0101: App
    "sequence_number": 0x0000,  # Set to 256 (0x0100) for options.
    "game_id": b"SDGS",
    "game_timestamp": {
        "year": 2023,
        "month": 12,
        "day": 28,
        "hour": 16,
        "minute": 34,
        "second": 43,
        "milli": 0,
    },
    "game_version": {
        "release": 0,
        "minor": 30,
        "major": 1,
    },
    "block_size": 0x40000,
    "header_block_count": 8,
    "unk1": 0,
    "hw_family": b"ACA",
    "hw_generation": 0,
    "org_timestamp": {
        "year": 0,
        "month": 0,
        "day": 0,
        "hour": 0,
        "minute": 0,
        "second": 0,
        "milli": 0,
    },
    "org_version": {
        "release": 0,
        "minor": 0,
        "major": 0,
    },
    "os_version": {
        "release": 1,
        "minor": 54,
        "major": 80,
    },
    "strings": b"\x00" * 0x27AC,  # Depending on the app/opt/pack, this might have some text in it.
}
# ----

# ---- Keys
# The BootID (app/opt/pack header) encryption key and IV.
BTKEY = bytes.fromhex("")
BTIV = bytes.fromhex("")

# The HMAC key that ensures the app/opt/pack created is authentic.
SIGKEY = bytes.fromhex("")
# ----

# ---- Constants. Don't edit.
EXFAT_HEADER = bytes.fromhex("EB769045584641542020200000000000")
OPTION_IV = bytes.fromhex("C063BF6F562D084D7963C987F5281761")
# ---- 

# ---- Script
Timestamp = Struct(
    "year" / Int16ul,
    "month" / Int8ul,
    "day" / Int8ul,
    "hour" / Int8ul,
    "minute" / Int8ul,
    "second" / Int8ul,
    "milli" / Int8ul,
)

Version = Struct(
    "release" / Int8ul,
    "minor" / Int8ul,
    "major" / Int16ul,
)

BootID = Struct(
    "length" / Const(0x2800, Int32ul),
    "magic" / Const(b"BTID", Bytes(4)),
    "type" / Int16ul,
    "sequence_number" / Int16ul,
    "game_id" / Bytes(4),
    "game_timestamp" / Timestamp,
    "game_version" / Version,
    "block_count" / Int64ul,
    "block_size" / Int64ul,
    "header_block_count" / Int64ul,
    "unk1" / Int64ul,
    "hw_family" / Bytes(3),
    "hw_generation" / Int8ul,
    "org_timestamp" / Timestamp,
    "org_version" / Version,
    "os_version" / Version,
    "strings" / Bytes(0x27AC),
)

def get_page_iv(iv: bytes, offset: int):
    return bytes(x ^ (offset >> (8 * (i % 8))) & 0xFF for (i, x) in enumerate(iv))


iv = secrets.token_bytes(16)

if BOOTID["type"] == 0x0201:
    iv = bytes(x ^ EXFAT_HEADER[i] ^ OPTION_IV[i] for (i, x) in enumerate(iv))

print(f"Generated IV: {iv.hex()}")

filesize = os.stat(INPUT_FILE).st_size
BOOTID["block_count"] = ceil(filesize / BOOTID["block_size"]) + 8

bullshit = secrets.token_bytes(BOOTID["block_size"])
bullshit_crc32 = zlib.crc32(bullshit)
block_crc32s = [0, bullshit_crc32, bullshit_crc32, bullshit_crc32, bullshit_crc32, bullshit_crc32, bullshit_crc32, bullshit_crc32]

with open(INPUT_FILE, "rb") as fin, open(OUTPUT_FILE, "w+b") as fout:
    # Write the bootID.
    cipher = AES.new(BTKEY, AES.MODE_CBC, BTIV)
    bootid = BootID.build(BOOTID)
    bootid_crc32 = zlib.crc32(bootid)
    bootid_bytes = cipher.encrypt(struct.pack("<I", bootid_crc32) + bootid)

    _ = fout.write(bootid_bytes)

    # Ignore the CRC32s and the signature for now, we'll come back later.
    # We'll generate random bytes for them though.
    _ = fout.write(secrets.token_bytes(BOOTID["block_size"] - 0x2800))

    # Bullshit out 7 random blocks for the header.
    for i in range(BOOTID["header_block_count"] - 1):
        _ = fout.write(bullshit)

    # Encrypt the contents of the file.
    total_written = 0
    to_read = filesize
    block_crc32 = 0

    while to_read > 0:
        page_iv = get_page_iv(iv, total_written)
        contents = fin.read(4096)
        contents_len = len(contents)

        to_read -= contents_len

        if contents_len < 4096:
            contents += b"\x00" * (4096 - contents_len)
        
        cipher = AES.new(ENCRYPTION_KEY, AES.MODE_CBC, page_iv)
        encrypted = cipher.encrypt(contents)
        total_written += fout.write(encrypted)
        block_crc32 = zlib.crc32(encrypted, block_crc32)

        if (total_written % BOOTID["block_size"]) == 0:
            block_crc32s.append(block_crc32)
            block_crc32 = 0
    
    # Pad with null bytes if we have an unfinished block.
    if (total_written % BOOTID["block_size"]) != 0:
        null_byte_count = BOOTID["block_size"] - (total_written % BOOTID["block_size"])
        
        while null_byte_count > 0:
            page_iv = get_page_iv(iv, total_written)
            cipher = AES.new(ENCRYPTION_KEY, AES.MODE_CBC, page_iv)
            encrypted = cipher.encrypt(b"\x00" * 4096)
            total_written += fout.write(encrypted)
            block_crc32 = zlib.crc32(encrypted, block_crc32)

            null_byte_count -= 4096
        
        block_crc32s.append(block_crc32)
        block_crc32 = 0

    # We can finally revisit the CRC32s.
    _ = fout.seek(0x2A04)

    for crc32 in block_crc32s[1:]:
        _ = fout.write(struct.pack("<I", crc32))

    _ = fout.seek(0)
    block_0_crc32 = zlib.crc32(fout.read(0x2800))

    # Skip the HMAC signature and the first CRC32, which we're trying to calculate.
    _ = fout.seek(0x204, os.SEEK_CUR)
    block_0_crc32 = zlib.crc32(fout.read(BOOTID["block_size"] - 0x2800 - 0x204), block_0_crc32)
    
    _ = fout.seek(0x2A00)
    _ = fout.write(struct.pack("<I", block_0_crc32))

    # Calculate the HMAC signature.
    _ = fout.seek(0x2A00)
    hmac = HMAC.new(SIGKEY, digestmod=SHA1)
    to_read = BOOTID["block_size"] * 8 - 0x2A00

    while to_read > 0:
        block = fout.read(min(0x1000, to_read))
        if len(block) == 0:
            raise Exception("Ran out of data making the signature.")

        _ = hmac.update(block)
        to_read -= len(block)
    
    _ = fout.seek(0x2800)
    _ = fout.write(hmac.digest())