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

nevermind. Got it working. Thanks

I'll preface this with the fact that I have no idea what I'm doing when it comes to python, but I couldn't get the original file to generate the "Canada" version of the OpenWRT file. So I used AI to strip out everything except what I need to do that. This may brick the AP completely, I don't have enough knowledge to confirm it's all good (maybe someone in the comments eventually will).

EDIT: I ended up back on the official firmware. Speeds were around 100Mbps download on OpenWRT while official is around 600Mbps. One weird thing is after the flash through the OpenWRT UI, I did have to flash again through the TP-Link UI because otherwise the AP would reset every reboot. Not sure what went wrong there, but good now. This was handy for getting back to official: https://argsnd.github.io/tp-link-stock-firmware-converter/

Anyways, the command I used with the below script is:

python3 patch-transplant-ca.py -i EAP245v3_ca_5.0.5_[20220323-rel68784]_up_signed.bin -f openwrt-24.10.4-ath79-generic-tplink_eap245-v3-squashfs-factory.bin -o factory-ca.bin

#!/usr/bin/python3
import argparse
import struct
import hashlib
import collections

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):
        self.parts[name] = data

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

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

    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 name, data in self.parts.items():
            table += template.format(name=name, offset=offset, size=len(data)).encode('ascii')
            offset += len(data)
        table += b'\0'
        buf = bytearray([0xff]*self.PART_TABLE_SIZE)
        buf[0:len(table)] = table
        return bytes(buf)

    def checksum(self):
        salt = bytes([
            0x7a, 0x2b, 0x15, 0xed, 0x9b, 0x98, 0x59, 0x6d,
            0xe5, 0x04, 0xab, 0x44, 0xac, 0x2a, 0x9f, 0x4e
        ])
        m = hashlib.md5()
        m.update(salt)
        m.update(self.description_data())
        m.update(self.part_table())
        for _, data in self.parts.items():
            m.update(data)
        return m.digest()

    def save(self, f):
        chk = self.checksum()
        data = self.description_data() + self.part_table()
        for _, d in self.parts.items():
            data += d
        data = chk + data
        data = struct.pack('>I', len(data) + 4) + data
        if self.padding:
            data += self.padding
        f.write(data)

    @classmethod
    def from_file(cls, f):
        def parse_ptn(line):
            parts = line.split()
            if len(parts) == 6 and parts[0] == 'fwup-ptn':
                return (parts[1], int(parts[3], 16), int(parts[5], 16))
            return None

        (fw_size,) = struct.unpack('>I', f.read(4))
        checksum = f.read(0x10)
        f.seek(cls.DESCRIPTION_OFFSET)
        (desc_len,) = struct.unpack('>I', f.read(4))
        desc = f.read(desc_len) if desc_len <= cls.DESCRIPTION_SIZE - 4 else None

        f.seek(cls.PART_TABLE_OFFSET)
        bulk = f.read(fw_size - cls.PART_TABLE_OFFSET)
        padding = f.read() or None

        fw = cls(desc, padding)

        pt_end = bulk.find(b'\0', 0)
        if pt_end > 0:
            lines = [l.strip() for l in bulk[:pt_end].decode('ascii').splitlines()]
            for line in lines:
                info = parse_ptn(line)
                if info:
                    name, base, size = info
                    fw.set_part(name, bulk[base:base+size])
        return fw


def main():
    parser = argparse.ArgumentParser(description='Transplant CA partition from TP-Link firmware to OpenWrt image.')
    parser.add_argument('-i', '--input', type=argparse.FileType('rb'), required=True,
                        help='Official TP-Link firmware with CA data')
    parser.add_argument('-f', '--factory', type=argparse.FileType('rb'), required=True,
                        help='OpenWrt factory image to patch')
    parser.add_argument('-o', '--output', type=argparse.FileType('wb'), required=True,
                        help='Output patched factory image')
    args = parser.parse_args()

    src = SafeloaderFirmware.from_file(args.input)
    dst = SafeloaderFirmware.from_file(args.factory)

    if 'product-info' not in src.parts:
        print("Error: 'product-info' partition not found in input firmware.")
        return

    print("Copying 'product-info' (CA region) from official firmware...")
    dst.set_part('product-info', src.get_part('product-info'))

    print("Saving patched image...")
    dst.save(args.output)
    print("Done. Patched firmware written successfully.")


if __name__ == '__main__':
    main()

I have successfully used the script to install OpenWRT on a 245v3 (US) edition. I patched mine with the CA version by mistake a while ago and that actually worked even though it wasn't recommended. I just patched from CA to CA this time and with your script patched the CA version into OpenWRT.

Thank you for making this amazing script.