safeloader patching script for OpenWrt

patch-safeloader.py

#!/usr/bin/python3

# This file allows modifying a a safeloader factory image.
# It can be used to transplant a partition from one image into another, or to update the
# firmware version info.
# The procedure to calculate the hash and the required salt are described at:
# https://github.com/openwrt/openwrt/blob/master/tools/firmware-utils/src/tplink-safeloader.c

import argparse
import binascii
import collections
import hashlib
import struct

class SafeloaderFirmware:
    DESCRIPTION_OFFSET = 0x14
    DESCRIPTION_SIZE = 0x1000
    PART_TABLE_OFFSET = 0x1014
    PART_TABLE_SIZE = 0x800

    def __init__(self, description=b'', padding=None):
        self.parts = collections.OrderedDict()
        self.description = description
        self.padding = padding

    def set_part(self, name, data):
        if name in self.parts:
            self.parts[name] = data
        else:
            self.parts[name] = data
            self.parts.move_to_end(name)

    def get_part(self, name):
        return self.parts[name]

    def part_table(self):
        template = 'fwup-ptn {name:s} base 0x{offset:05x} size 0x{size:05x}\t\r\n'
        offset = self.PART_TABLE_SIZE
        table = b''
        for part_name,part_data in self.parts.items():
            table += template.format(name=part_name, offset=offset, size=len(part_data)).encode('ascii')
            offset += len(part_data)
        table += b'\0'
        data = bytearray([0xff]*self.PART_TABLE_SIZE)
        data[0:len(table)] = table
        return bytes(data)

    def description_data(self):
        data = bytearray([0xff]*self.DESCRIPTION_SIZE)
        desc = self.description
        if desc is not None:
            data[0:4] = struct.pack('>I', len(desc))
            data[4:4+len(desc)] = desc
        return bytes(data)

    def checksum(self):
        # The calculated hash is seeded with a salt, followed by the FW image body
        m = hashlib.md5()
        m.update(bytes([
            0x7a, 0x2b, 0x15, 0xed, 0x9b, 0x98, 0x59, 0x6d,
            0xe5, 0x04, 0xab, 0x44, 0xac, 0x2a, 0x9f, 0x4e
        ]))
        m.update(self.description_data())
        m.update(self.part_table())
        for name,part in self.parts.items():
            m.update(part)
        return m.digest()

    def save_firmware(self, output_file):
        checksum = self.checksum()

        data = self.description_data()
        data += self.part_table()
        for name,part in self.parts.items():
            data += part
        data = self.checksum() + data
        data = struct.pack('>I', len(data) + 4) + data

        if self.padding is not None:
            data += self.padding
        output_file.write(data)

    @staticmethod
    def unpack_metadata_partition(part):
        return (int.from_bytes(part[0:4], 'big'),  part[8:])

    @staticmethod
    def pack_metadata_partition(data):
        return len(data).to_bytes(4, 'big') + int(0).to_bytes(4) + data

    def set_version(self, new_version):
        part_soft_version = self.get_part('soft-version')
        part_len, pdata = self.unpack_metadata_partition(part_soft_version)

        if pdata.isascii():
            # Must add NULL termination
            pdata = (new_version + '\0').encode('ascii')
        elif part_len < 12:
            print("Cannot parse partition as structured version info")
            return
        else:
            v_maj, v_min, v_patch = new_version.split('.', 3)
            if not (v_maj.isdigit() and v_min.isdigit() and v_patch.isdigit()):
                print(f'failed to parse new version "{new_version}" as "MAJ.MIN.PATCH"')
                return

            pdata = bytearray(pdata)
            pdata[1] = int(v_maj)
            pdata[2] = int(v_min)
            pdata[3] = int(v_patch)

        self.set_part('soft-version', self.pack_metadata_partition(pdata))

    def set_compatibility(self, compatibility):
        part_soft_version = self.get_part('soft-version')
        part_len, pdata = self.unpack_metadata_partition(part_soft_version)

        if pdata.isascii() or part_len < 16:
            print("Cannot parse partition as structured version info")
            return

        pdata = bytearray(pdata)
        pdata[12:16] = compatibility.to_bytes(4, 'big')

        self.set_part('soft-version', self.pack_metadata_partition(pdata))

    @classmethod
    def from_file(cls, input_file):
        def read_fw_ptn_list(fw):
            fwup_end = fw.find(b'\0', 0)
            if fwup_end > 0:
                return [p.strip() for p in fw[:fwup_end].decode('ascii').splitlines()]
            else:
                return None
        
        def parse_ptn(ptn):
            tokens = ptn.split()
            if len(tokens) != 6:
                return None
            elif tokens[0] != 'fwup-ptn' or tokens[2] != 'base' or tokens[4] != 'size':
                return None
            else:
                return (tokens[1], int(tokens[3], base=16), int(tokens[5], base=16))
        
        (fw_size,) = struct.unpack('>I', input_file.read(4))
        checksum = input_file.read(0x10)
        input_file.seek(cls.DESCRIPTION_OFFSET)

        (desc_size,) = struct.unpack('>I', input_file.read(4))
        if desc_size <= cls.DESCRIPTION_SIZE - 4:
            desc = input_file.read(desc_size)
        else:
            desc = None

        input_file.seek(cls.PART_TABLE_OFFSET)
        bulk = input_file.read(fw_size-cls.PART_TABLE_OFFSET)
        padding = input_file.read()
        if len(padding) == 0:
            padding = None

        fw = cls(desc, padding)

        part_list = read_fw_ptn_list(bulk)
        for ptn in part_list:
            name, base, size = parse_ptn(ptn)
            fw.set_part(name, bulk[base:base+size])

        if fw.checksum() != checksum:
            print('invalid file checksum')
        return fw


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Patch OpenWrt safeloader factory image')
    parser.add_argument('-f', '--factory',
            type=argparse.FileType('rb'),
            required=True,
            help='OpenWrt factory image to be patched')
    parser.add_argument('-o', '--output',
            type=argparse.FileType('wb'),
            help='output path for the patched file')

    subparsers = parser.add_subparsers(dest='mode')
    parser_transplant = subparsers.add_parser('transplant', help='copy a partition')
    parser_transplant.add_argument('-i', '--input',
            type=argparse.FileType('rb'),
            required=True,
            help='Safeloader image to read partition from')
    parser_transplant.add_argument('-p', '--partition',
            type=str,
            required=True,
            help='Partition to patch into OpenWrt factory image')

    parser_version = subparsers.add_parser('version', help='modify firmware version numbers')
    parser_version.add_argument('-v', '--version',
            type=str,
            help="New firmware version")
    parser_version.add_argument('-c', '--compatibility',
            type=int,
            help="New compatibility level (if supported)")

    args = parser.parse_args()

    # Read OpenWrt factory image
    factory = SafeloaderFirmware.from_file(args.factory)

    if args.mode == 'transplant':
        # Read patch source input file
        patch_input = SafeloaderFirmware.from_file(args.input)

        if args.patch is not None:
            if args.patch in patch_input.parts:
                print(f'Patching "{args.patch}" into OpenWrt factory image...')
                factory.set_part(args.patch, patch_input.get_part(args.patch))
            else:
                print(f'Could not find firmware part "{args.patch}" in input image')
    elif args.mode == 'version':
        if args.version is not None:
            factory.set_version(args.version)
        if args.compatibility is not None:
            factory.set_compatibility(args.compatibility)

    # Write updated factory image
    if args.output:
        factory.save_firmware(args.output)
    else:
        print('Patched factory image not saved, resulting partition table:')
        for name, data in factory.parts.items():
            print(f'\t{name:s} +0x{len(data):06x}')

@kaelbashir @svanheule Great news... Searched different keywords, came across r/TPLinkOmada/Safe to flash a US FW onto a Canadian AP?... Flashing via TFTP bypasses the region check. Embarrassed I didn't think of this before, though not as embarrassed as TP-Link should be for not publishing an update containing security patches over 6 months old, for every region except the US...