APT:Wrap/APT:Unwrap reimplementation in python

apt_wraps.py

from Cryptodome.Cipher import AES

wrap_key = ... # keyslot 0x31 normalkey, obtain KeyX from boot9 and KeyY from NATIVE_FIRM, use something like pyctr to keygen

def apt_unwrap(inp: bytes, nnc_offset: int, nnc_size: int, skip_verify_nonstandard_ccm: bool = False):
    unwrapped_len = len(inp) - 0x10

    nnc_size &= 12
    nonce = inp[0:nnc_size] + b'\x00' * (12 - nnc_size)

    data_after_nonce_size = unwrapped_len - (nnc_offset + nnc_size)
    total_ccm_crypt_size = nnc_offset + data_after_nonce_size
    mac_offset = nnc_size + total_ccm_crypt_size

    mac = inp[mac_offset:mac_offset+0x10]

    dec = AES.new(key=wrap_key, mode=AES.MODE_CCM, nonce=nonce, mac_len=16) \
        .decrypt(inp[nnc_size:nnc_size+total_ccm_crypt_size])

    if not skip_verify_nonstandard_ccm:
        _, testmac = AES.new(key=wrap_key, mode=AES.MODE_CCM, nonce=nonce, mac_len=16) \
            .encrypt_and_digest(dec + b'\x00' * (len(dec) % 0x10))

        if testmac != mac:
            raise Exception("MAC mismatch")

    return dec[:nnc_offset] + nonce[:nnc_size] + dec[nnc_offset:]

def apt_wrap(inp: bytes, nnc_offset: int, nnc_size: int):
    nnc_size &= 12
    nonce = inp[nnc_offset:nnc_offset+nnc_size] + b'\x00' * (12 - nnc_size)

    encbuf = inp[0:nnc_offset] + inp[nnc_offset+nnc_size:]
    padlen = len(encbuf) % 0x10
    encbuf += b'\x00' * padlen

    encrypted, outmac = AES.new(key=wrap_key, mode=AES.MODE_CCM, nonce=nonce, mac_len=16) \
        .encrypt_and_digest(encbuf)

    return nonce[:nnc_size] + encrypted[:-padlen] + outmac