Python script to convert IMG4 to SHSH blobs

convert.py

"""
Converts a raw img4 to an SHSH blob.
Similar to `img4tool --convert -s out.shsh dump.raw`

Requirements: 
  pyasn1

Usage: 
  python convert.py <raw img4> [output]
"""
import argparse
import plistlib
import pyasn1.codec.der.decoder
import pyasn1.codec.der.encoder
import pyasn1.type.univ
from typing import Union


def get_im4m_from_img4(img4: pyasn1.type.univ.Sequence) -> pyasn1.type.univ.Sequence:
    im4m = img4[2]
    assert str(im4m[0]) == "IM4M"
    assert isinstance(im4m, pyasn1.type.univ.Sequence)
    return im4m


def get_im4r_from_img4(img4: pyasn1.type.univ.Sequence) -> pyasn1.type.univ.Sequence:
    im4r = img4[3]
    assert str(im4r[0]) == "IM4R"
    assert isinstance(im4r, pyasn1.type.univ.Sequence)
    return im4r


def get_bncn_from_im4r(im4r: pyasn1.type.univ.Sequence) -> pyasn1.type.univ.Sequence:
    bncn = im4r[-1][-1]
    assert str(bncn[0]) == "BNCN"
    assert isinstance(bncn, pyasn1.type.univ.Sequence)
    return bncn


def endian_converter(val: Union[bytes, str, pyasn1.type.univ.OctetString]) -> str:
    """Converts big endian to small endian and vice versa

    Args:
        val (Union[bytes, str]): Bytes or a hex string to convert

    Returns:
        str: A 0x prefixed hex string
    """
    ba = bytearray()
    if isinstance(val, str):
        ba = bytearray.fromhex(str.removeprefix("0x"))
    else:
        ba = bytearray(val)
    ba.reverse()
    s = ''.join(format(x, '02x') for x in ba)
    return f'0x{s}'


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('file', help="File to convert")
    parser.add_argument('out', nargs="?", default="out.shsh", help="Output file")
    
    args = parser.parse_args()
    img4_file = args.file
    with open(img4_file, 'rb') as f:
        img4, _ = pyasn1.codec.der.decoder.decode(f.read())
    im4r = get_im4r_from_img4(img4)
    bncn = get_bncn_from_im4r(im4r)
    generator = endian_converter(bncn[-1])
    print(f"Generator: {generator}")

    # Weird fixup
    im4m = get_im4m_from_img4(img4)
    im4m.tagSet._TagSet__superTags = (im4m.tagSet._TagSet__superTags[0],)

    shsh = {}
    shsh["ApImg4Ticket"] = pyasn1.codec.der.encoder.encode(im4m)
    shsh["generator"] = generator
    with open(args.out, 'wb') as f:
        plistlib.dump(shsh, f)
        print(f"Saved blob to {args.out}")


if __name__ == "__main__":
    main()