Convert TP-Link Switch Firwmare Upgrade to sysupgrade file

switchfw2sysupgrade.py

#!/usr/bin/env python

from optparse import OptionParser
import tempfile
import subprocess
import os
import shutil
import sys
import binascii
import struct

TP_IMAGE_HEADER_SIZE            = 512
MD5_DIGEST_LEN                  = 16
FLASH_PRODUCT_LENGTH            = 16
TP_IMAGE_NAME_SIZE              = 200
USER_IMAGE_SIZE                 = 4
UIMAGE_SIZE                     = 4
BOOTROM_SIZE                    = 4
COMPRESS_METHOD                 = 4
FL_SIZE_LEN                     = 16
KERNEL_SIZE                     = 6291456

def decrypt_image(infile, directory):
    outfile = directory.name + os.path.sep + "decrypted.tmp"
    ret = subprocess.run(["openssl", "enc", "-d", "-des-cbc", "-in", infile, "-out", outfile, "-nosalt", "-nopad", "-p", "-iv", "f51010736efbabb2", "-K", "0001020304050607"], stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
    if ret.returncode != 0:
        print("Error decrypting image.")
        sys.exit(1)
    return outfile

def decode_header(fn):
    header_info = {}
    fp = open(fn, 'rb')
    header = fp.read(TP_IMAGE_HEADER_SIZE)
    fp.close()
    header_info['productId'] = header[MD5_DIGEST_LEN: MD5_DIGEST_LEN + FLASH_PRODUCT_LENGTH]
    header_info['imagename'] = header[MD5_DIGEST_LEN + FLASH_PRODUCT_LENGTH : MD5_DIGEST_LEN + FLASH_PRODUCT_LENGTH + TP_IMAGE_NAME_SIZE]
    offset = MD5_DIGEST_LEN + FLASH_PRODUCT_LENGTH + TP_IMAGE_NAME_SIZE
    header_info['userImageSize'] = struct.unpack(">I", header[offset : offset + USER_IMAGE_SIZE])[0]
    offset += USER_IMAGE_SIZE
    header_info['uImageSize'] = struct.unpack(">I", header[offset : offset + UIMAGE_SIZE])[0]
    offset += UIMAGE_SIZE
    header_info['bootromSize'] = struct.unpack(">I", header[offset : offset + BOOTROM_SIZE])[0]
    offset += BOOTROM_SIZE
    header_info['compressMethod'] = struct.unpack(">I", header[offset : offset + COMPRESS_METHOD])[0]
    offset += COMPRESS_METHOD
    for size in ['flBootSize', 'flKernelOffset', 'flKernelSize', 'flUsrImg1Offset', 'flUsrImg1Size', 'flUsrImg2Offset', 'flUsrImg2Size', 'flUsrAppOffset', 'flUsrAppSize', 'flVer']:
        header_info[size] = header[offset : offset + FL_SIZE_LEN]
        offset += FL_SIZE_LEN

    print("Image Name     : " + header_info['imagename'].strip(b'\x00').decode("ascii"))
    print("User Image Size: " + hex(header_info['userImageSize']))
    print("uImage Size    : " + hex(header_info['uImageSize']))
    print("bootromSize    : " + hex(header_info['bootromSize']))
    print("compressMethod : " + hex(header_info['compressMethod']))
    print("flBootSize     : " + header_info['flBootSize'].strip(b'\x00').decode("ascii"))
    print("flKernelOffset : " + header_info['flKernelOffset'].strip(b'\x00').decode("ascii"))
    print("flKernelSize   : " + header_info['flKernelSize'].strip(b'\x00').decode("ascii"))
    print("flUsrImg1Offset: " + header_info['flUsrImg1Offset'].strip(b'\x00').decode("ascii"))
    print("flUsrImg1Size  : " + header_info['flUsrImg1Size'].strip(b'\x00').decode("ascii"))
    return header_info

def extract_part(fn, offset, length, out_file):
    fp = open(out_file, 'wb')
    img = open(fn, 'rb')
    img.seek(offset)
    fp.write(img.read(length))
    img.close()
    fp.close()

def extract_user_image(fn, info, temp_dir):
    offset = TP_IMAGE_HEADER_SIZE
    length = info['userImageSize']
    out_file = temp_dir.name + os.path.sep + "user_image.bin"
    extract_part(fn, offset, length, out_file)
    return out_file

def extract_kernel(fn, info, temp_dir):
    offset = TP_IMAGE_HEADER_SIZE + info['userImageSize']
    length = info['uImageSize']
    out_file = temp_dir.name + os.path.sep + "kernel.bin"
    extract_part(fn, offset, length, out_file)
    return out_file

if __name__ == "__main__":
    parser = OptionParser()
    parser.add_option("-o", "--output", dest="output", help="write sysupgrade image to FILE", metavar="FILE")
    parser.add_option("-i", "--input", dest="input", help="read encrypted firmware image from FILE", metavar="FILE")
    
    (options, args) = parser.parse_args()
    if not options.input or not options.output:
        parser.error("Input and Output files are required")

    if not shutil.which("openssl"):
        print("OpenSSL is required in PATH to run this script")
        sys.exit(1)

    temp_dir = tempfile.TemporaryDirectory()
    decrypted = decrypt_image(options.input, temp_dir)
    header_info = decode_header(decrypted)
    kernel = extract_kernel(decrypted, header_info, temp_dir)
    user_image = extract_user_image(decrypted, header_info, temp_dir)
    out_file = open(options.output, 'wb')
    fp = open(kernel, 'rb')
    out_file.write(fp.read())
    fp.close()
    out_file.truncate(KERNEL_SIZE)
    out_file.seek(KERNEL_SIZE)
    fp = open(user_image, 'rb')
    out_file.write(fp.read())
    out_file.close()
    
    temp_dir.cleanup()