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}')

Running it directly on the TPLink firmware gives an invalid file checksum error (tested on v3.0 and v3.1, on both US and Canadian firmware). Attempting to update the AP with the output patched-factory.bin also fails.

$ python patch-safeloader.py -f openwrt-factory.bin -i EAP235-30.bin  -p soft-version -o patched.bin
invalid file checksum
Patching "soft-version" into OpenWrt factory image...

I get the same error message:

$ bin/patch-safeloader -f src/openwrt-firmware-eap235wall.bin -i src/official-firmware-3p1p0-eap235wall.bin -p soft-version -o out/openwrt-firmware-eap235wall-patched.bin
invalid file checksum
Patching "soft-version" into OpenWrt factory image...

The successfull version mirror (https://svanheule.net/files/eap235-wall-v1-squashfs-factory-v3-test.bin) also seems to be down, any chance this could be hosted on openwrt in the meanwhile?

Thanks in advance!

@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...