bpatch.py

#!/usr/bin/python3

"""bpatch.py - Bit Patch Firmware Tool"""

__author__ = "Andrea De Simone and Fabrizio Riente"
__copyright__ = "Copyright 2025, Politecnico di Torino"
__license__ = "Apache-2.0"
__version__ = "2.0.0"
__maintainer__ = "Andrea De Simone"
__email__ = "andrea.desimone@polito.it, fabrizio.riente@polito.it"

import sys
import os
from math import ceil
import shutil

# private import
from path_bpatch import path_patch_exec
from binary_generation import write_binary_patch, read_binary_patch, generate_hex_fw
from delta_generation import decode_py, encode_fw

# enable custom header
header_custom = False
# from custom_header.header_sbsfu import header_fw_size, header_patch_size, header_lines, write_header_custom

""" Custom Header definition start """
if not header_custom:
    header_fw_size    = 0   # size of the custom header in the firmware in bytes
    header_patch_size = 0   # size of the custom header in the patch in bits
    header_lines      = 0   # number of lines of the custom header in the patch

    write_header_custom = None  # custom header disabled
"""
# function to build header for txt patch
def write_header_custom(bin_txt):
    """"""Read the header from "new_fw_header.tmp" and write the custom header in the patch in binary
    
    :param bin_txt: string with the binary patch
    :return: string with the binary patch with the header added
    """"""
    pass
"""

""" Custom Header definition end """


if __name__ == '__main__':
    if len(sys.argv) < 2 or sys.argv[1] not in ['encode', 'decode', 'help']:
        print("Usage: python3 patch.py [OPTION]")
        print("Options:"
              "\n\t-encode: generate the patch"
              "\n\t-decode: apply the patch"
              "\n\t-help  : help")
        exit()

    # get pid for temporary files
    PID = os.getpid()

    # encode function
    if sys.argv[1] == 'encode':
        # check the number of arguments
        if len(sys.argv) < 5:
            print("Usage: python3 patch.py encode OLD_FW NEW_FW PATCH [OPTIONS]")
            print("Options:"
                  "\n\t-d         : use diff with option minimal to reduce the patch size"
                  "\n\t-e         : use encode dict opcodes to reduce the patch size"
                  "\n\t-v         : verify che correctness of the patch"
                  "\n\t-V         : verify che correctness of the txt patch"
                  "\n\t-t FILENAME: write the patch in txt format"
                  "\n\t-b FILENAME: write the patch in binary format"
                  "\n\t-r FILENAME: write the report in txt format"
                  "\n\t-R FILENAME: write the report in csv format (if the file exists with correct header, append the report)")
            exit()

        # decode arguments
        old_fw = sys.argv[2]
        new_fw = sys.argv[3]
        p = sys.argv[4]

        # default values
        verify = False
        verify_py = False
        minimal = False
        encode_dict = False
        path_patch_txt = ""
        path_patch_bin_txt = ""
        path_report_txt = ""
        path_report_csv = ""
        # check the optional arguments
        try:
            for i in range(5, len(sys.argv)):
                # verify the correctness of the patch
                if sys.argv[i] == '-v':
                    verify = True
                # verify the flow of the patch generation
                elif sys.argv[i] == '-V':
                    verify_py = True
                # write the patch in txt format
                elif sys.argv[i] == '-t':
                    path_patch_txt = sys.argv[i + 1]
                # write the patch in binary format
                elif sys.argv[i] == '-b':
                    path_patch_bin_txt = sys.argv[i + 1]
                # write report in txt format
                elif sys.argv[i] == '-r':
                    path_report_txt = sys.argv[i + 1]
                # write report in csv format
                elif sys.argv[i] == '-R':
                    path_report_csv = sys.argv[i + 1]
                # use minimal option in diff
                elif sys.argv[i] == '-d':
                    minimal = True
                # enable encode dict
                elif sys.argv[i] == '-e':
                    encode_dict = True

        except IndexError:
            print("Error during parsing arguments")

        os.makedirs(f"{PID}", exist_ok=True)

        if verify_py:
            if not path_report_txt:
                path_patch_txt = f"{PID}/patch_txt.tmp"
            if not path_patch_bin_txt:
                path_patch_bin_txt = f"{PID}/patch_bin_txt.tmp"

        # read the binary files
        generate_hex_fw(old_fw, f"{PID}/old_fw.tmp", f"{PID}/old_fw_header.tmp", header_fw_size)
        generate_hex_fw(new_fw, f"{PID}/new_fw.tmp", f"{PID}/new_fw_header.tmp", header_fw_size)
        # compute diff
        if minimal:
            os.system(f"diff -d {PID}/old_fw.tmp {PID}/new_fw.tmp > {PID}/diff.tmp")
        else:
            os.system(f"diff {PID}/old_fw.tmp {PID}/new_fw.tmp > {PID}/diff.tmp")
        # parse diff
        os.system(f"grep -E '^[0-9,]+[acd][0-9,]+$' {PID}/diff.tmp > {PID}/diff_c.tmp")
        os.system(f"grep '^>' {PID}/diff.tmp | sed 's/^> //' > {PID}/diff_a.tmp")

        # get the size of the firmware
        size_old_fw = os.path.getsize(old_fw)
        size_new_fw = os.path.getsize(new_fw)

        # encode function
        wnbd, wnbc, wnba, patch_size, patch, dict_add, dict_cpy = encode_fw(f"{PID}/diff_c.tmp", f"{PID}/diff_a.tmp", size_old_fw - header_fw_size, size_new_fw, custom_header_size=header_patch_size, path_report_txt=path_report_txt, path_report_csv=path_report_csv, encode_dict=encode_dict)

        # write the txt patch
        if path_patch_txt:
            with open(path_patch_txt, "w") as f:
                f.write(patch)

        # write binary patch
        write_binary_patch(patch, p, wnbd, wnbc, wnba, patch_size, dict_add, dict_cpy, write_header_custom, path_patch_bin_txt)

        # verify that the patch_size is correct
        actual_patch_size_bytes = os.path.getsize(p)
        if ceil(patch_size/8) != actual_patch_size_bytes:
            print(f"Warning: the computed patch size {ceil(patch_size/8)} bytes is different from the actual patch size {actual_patch_size_bytes} bytes")

        if verify_py:
            # verify the correctness of the txt patch
            decode_py(f"{PID}/old_fw.tmp", path_patch_txt, f"{PID}/patched_fw.tmp")

            # check if the patched firmware is equal to the new firmware
            if os.system(f"diff {PID}/patched_fw.tmp {PID}/new_fw.tmp -q"):
                print("The txt patch is not correct")
                shutil.rmtree(f"{PID}")
                exit(2)
            else:
                print("The txt patch is correct")

            # verify the correctness of the binary patch
            read_binary_patch(path_patch_bin_txt, f"{PID}/patch_txt_new.tmp", header_lines)

            # check if the patched firmware is equal to the new firmware
            if os.system(f"diff {path_patch_txt} {PID}/patch_txt_new.tmp -q"):
                print("The bin patch is not correct")
                shutil.rmtree(f"{PID}")
                exit(3)
            else:
                print("The bin patch is correct")

        # check the correctness of the patch
        if verify or verify_py:
            # generate the patched firmware
            os.system(f"{path_patch_exec} {old_fw} {p} {PID}/patched_fw.tmp {1016} {1024} > /dev/null")
            # check if the patched firmware is equal to the new firmware
            if os.system(f"diff {PID}/patched_fw.tmp {new_fw} -q"):
                print("The patch is not correct")
                shutil.rmtree(f"{PID}")
                exit(1)
            else:
                print("The patch is correct")

        # remove temporary files
        shutil.rmtree(f"{PID}")

        exit(0)

    # decode function
    elif sys.argv[1] == 'decode':
        # check the number of arguments
        if len(sys.argv) < 5:
            print("Usage: python3 patch.py decode OLD_FW PATCH NEW_FW [OPTIONS]")
            print("Options:"
                  "\n\t-v     : enable verbose mode"
                  "\n\t-r SIZE: set the read buffer size"
                  "\n\t-p SIZE: set the patch buffer size")
            exit()

        old_fw = sys.argv[2]
        p = sys.argv[3]
        new_fw = sys.argv[4]

        # default values
        verbose =  '> /dev/null'
        read_buffer = 1016
        patch_buffer = 1024
        # check the optional arguments
        try:
            for i in range(5, len(sys.argv)):
                # enable verbose mode
                if sys.argv[i] == '-v':
                    verbose = ''
                # set the read buffer size
                elif sys.argv[i] == '-r':
                    read_buffer = int(sys.argv[i + 1])
                # set the patch buffer size
                elif sys.argv[i] == '-p':
                    patch_buffer = int(sys.argv[i + 1])

        except IndexError:
            print("Error during parsing arguments")

        # apply the patch
        exit(os.system(f"{path_patch_exec} {old_fw} {p} {new_fw} {read_buffer} {patch_buffer} {verbose}"))

    # help function
    elif sys.argv[1] == 'help':
        print("Usage: python3 patch.py [OPTION]")
        print("\n\tOptions:"
              "\n\t-encode: generate the patch"
              "\n\t\t Usage: python3 patch.py encode OLD_FW NEW_FW PATCH [OPTIONS]")
        print("\n\t\tOptions:"
              "\n\t\t-d         : use diff with option minimal to reduce the patch size"
              "\n\t\t-e         : use encode dict opcodes to reduce the patch size"
              "\n\t\t-v         : verify che correctness of the patch"
              "\n\t\t-V         : verify che correctness of the txt patch"
              "\n\t\t-t FILENAME: write the patch in txt format"
              "\n\t\t-b FILENAME: write the patch in binary format"
              "\n\t\t-r FILENAME: write the report in txt format"
              "\n\t\t-R FILENAME: write the report in csv format")

        print("\n\t-decode: apply the patch"
              "\n\t\t Usage: python3 patch.py decode OLD_FW PATCH NEW_FW [OPTIONS]")
        print("\n\t\tOptions:"
              "\n\t\t-v     : enable verbose mode"
              "\n\t\t-r SIZE: set the read buffer size"
              "\n\t\t-p SIZE: set the patch buffer size")

        print("\n\t-help")
        exit()

    else:
        print("Usage: python3 patch.py [OPTION]")
        print("Options:"
              "\n\t-encode: generate the patch"
              "\n\t-decode: apply the patch"
              "\n\t-help  : help")
        exit()