Rewritten Python script to split Insyde H2O firmware files

Question

I’ve recently tried to update my Lenovo laptop’s firmware without installing Windows. I haven’t done it yet, but during their research someone on Matrix linked me to this page (original source, archived) which contained a Python script that splits firmware files. I didn’t want to run code w/o reviewing it first and I didn’t like their implementation. I’ve written my own version of the script which I would like feedback on. It works both as a module and as an executable.

#!/usr/bin/env python3
from dataclasses import dataclass

__all__ = ["default_splitter", "Splitter", "Component"]


@dataclass
class Component:
    magic: bytes
    name: str


class Splitter:
    def __init__(self, components, header_name, tag_len, offset_dist):
        self.components = components
        self.header_name = header_name
        self.tag_len = tag_len
        self.offset_dist = offset_dist

    def open(self, content):
        if isinstance(content, bytes):
            self.content = content
        else:
            obj = open(content, "rb")
            self.ctx = obj
            self.content = obj.read()
            return obj

    def close(self):
        assert self.ctx is not None
        self.ctx.close()

    def find(self):
        assert self.content is not None
        self.positions = {}
        for component in self.components:
            self.positions[component.name] = self.content.find(component.magic)

    def position_of(self, component):
        return self.positions[component.name]

    def get_slice(self, begin, end):
        return self.content[begin:end]

    def contents(self):
        assert self.content is not None
        assert self.positions is not None
        yield self.header_name, self.get_slice(
            0, self.position_of(self.components[0]) if bool(self.components) else None
        )
        for component, next_component in lookahead(self.components):
            begin = self.position_of(component) + self.tag_len + self.offset_dist
            assert begin is not None
            end = (
                self.position_of(next_component) if next_component is not None else None
            )
            yield component.name, self.content[begin:end]
            if end is None:
                break


default_splitter = Splitter(
    [
        Component(b"$_IFLASH_DRV_IMG", "drv_img"),
        Component(b"$_IFLASH_BIOSIMG", "biosimg"),
        Component(b"$_IFLASH_INI_IMG", "ini_img"),
        Component(b'$_IFLASH_EC_IMG_', "ec_img")
    ],
    "head",
    16,
    8,
)


# adapted from: https://stackoverflow.com/a/1630350/8025936
def lookahead(iterable):
    iterator = iter(iterable)
    next_item = next(iterator)
    for item in iterator:
        yield next_item, item
        next_item = item
    yield next_item, None


if __name__ == "__main__":
    from argparse import ArgumentParser
    from os import makedirs
    from os.path import exists
    from pathlib import Path

    parser = ArgumentParser(description="Split Insyde H2O firmware files")
    parser.add_argument(
        "-f",
        "--file",
        dest="filename",
        help="Input file name",
        metavar="file",
        type=Path,
    )
    parser.add_argument(
        "-o",
        "--dest",
        dest="dest_dir",
        help="Destination directory",
        metavar="destination",
        type=Path,
    )
    args = parser.parse_args()
    if not exists(args.dest_dir):
        makedirs(args.dest_dir)
    with default_splitter.open(args.filename):
        default_splitter.find()
        for name, content in default_splitter.contents():
            dest = (args.dest_dir / name).with_suffix(".bin")
            with open(dest, "wb") as f:
                f.write(content)

Answer 1

1. Make you api easy to use and hard to misuse. Right now, your Splitter class requires a lot of initialization to be used. You have to first call open, then find and then call contents. The job of the splitter is to split. Ideally you should just have one method called split, instead of relying on a bunch of asserts that don't even tell the user what it failed.
2. Use __function_name for private methods. This adds clarity to your api.
3. Maintain seperation of concerns. The splitter should not be responsible for opening and reading file content.
4. Use assertions to check preconditions. Scattering assertions in between logic makes the code less readable. E.g. assert begin is not None makes no sense. You define begin = self.position_of(component) + self.tag_len + self.offset_dist. There is no way begin can be None.
5. Dont abstract out built-in syntax. get_slice method is perfectly define in python using the [begin:end] syntax. Someone familar with python will immediately understand this whereas for get_slice you have to go look at the implementation.
6. Similarly for position_of. Looking at it for the first time, I thought position_of must be some firmware specific operation to find position of a component. But turns out it just finding the position of the component in the dict of positions u created.
7. There is a possibility that .find() might return -1. I dont see any checks for that.
8. Instead of defining end conditionally. You can just check if next_component is None and branch. It makes the code readable.

#!/usr/bin/env python3
from dataclasses import dataclass

__all__ = ["default_splitter", "Splitter", "Component"]


@dataclass
class Component:
    magic: bytes
    name: str


class Splitter:
    def __init__(self, components, header_name, tag_len, offset_dist):
        self.components = components
        self.header_name = header_name
        self.tag_len = tag_len
        self.offset_dist = offset_dist


    def split(self, content):

        assert isinstance(content, bytes), "Content must be a byte string"

        positions = {positions[component.name] : content.find(component.magic)
            for component in self.components}

        assert positions, "Content does not contain any component"

        yield self.header_name, self.content[0 :  positions[component[0].name]]

        for component, next_component in lookahead(self.components):

            begin = positions[component.name] + self.tag_len + self.offset_dist

            if next_component is None:
                yield component.name, content[begin:]
                break
            
            yield component.name, content[begin:positions[next_component.name]]
                

default_splitter = Splitter(
    [
        Component(b"$_IFLASH_DRV_IMG", "drv_img"),
        Component(b"$_IFLASH_BIOSIMG", "biosimg"),
        Component(b"$_IFLASH_INI_IMG", "ini_img"),
        Component(b'$_IFLASH_EC_IMG_', "ec_img")
    ],
    "head",
    16,
    8,
)


# adapted from: https://stackoverflow.com/a/1630350/8025936
def lookahead(iterable):
    iterator = iter(iterable)
    next_item = next(iterator)
    for item in iterator:
        yield next_item, item
        next_item = item
    yield next_item, None


if __name__ == "__main__":
    from argparse import ArgumentParser
    from os import makedirs
    from os.path import exists
    from pathlib import Path

    parser = ArgumentParser(description="Split Insyde H2O firmware files")
    parser.add_argument(
        "-f",
        "--file",
        dest="filename",
        help="Input file name",
        metavar="file",
        type=Path,
    )
    parser.add_argument(
        "-o",
        "--dest",
        dest="dest_dir",
        help="Destination directory",
        metavar="destination",
        type=Path,
    )
    args = parser.parse_args()
    if not exists(args.dest_dir):
        makedirs(args.dest_dir)
    with open(args.filename, 'rb') as input_file:
        for name, content in default_splitter.split(input_file.read()):
            dest = (args.dest_dir / name).with_suffix(".bin")
            with open(dest, "wb") as output_file:
                output_file.write(content)