ALT1160 flash image splitter

splitter.py

#!/usr/bin/env python3

from collections import namedtuple
from re import compile as regex
import sys

Part = namedtuple('Part', 'name, size, offset')


class Flashdump(object):
    def __init__(self, mtdparts):
        self._mtdparts = mtdparts

    def __repr__(self):
        return "\n".join(["{:<20s} 0x{:0>8x} 0x{:0>8x}".format(p.name.decode(), p.size, p.offset) for p in self._mtdparts])

    def split(self, file):
        with open(file, "br") as fin:
            for p in self._mtdparts:
                with open(p.name, "bw") as fout:
                    fin.seek(p.offset)
                    fout.write(fin.read(p.size))


class Uboot(object):
    def __init__(self, file):
        self._file = file
        self._mtdparts = None

    def each_chunk(self, stream, separator):
        buffer = bytearray()
        while True:  # until EOF
            chunk = stream.read(4096)
            if not chunk:  # EOF?
                yield buffer
                break
            buffer.extend(chunk)
            while True:  # until no separator is found
                try:
                    part, buffer = buffer.split(separator, 1)
                except ValueError:
                    break
                else:
                    yield part

    def parse(self, line):
        pattern = regex(b"^(\d+)([km]?)\((\w+)\)$")
        unit = {b"k": 1024, b"m": 1024 * 1024}

        mtdparts = list()
        offset = 0
        for p in line.split(b","):
            m = pattern.match(p)
            if m:
                size = int(m.group(1)) * unit[m.group(2)]
                name = m.group(3)
                mtdparts.append(Part(name, size, offset))
                offset += size
            else:  # last partition -(tstorage)
                mtdparts.append(Part(b"tstorage", -1, offset))
                print("last partition", p)
                break

        self._mtdparts = mtdparts

    def mtdparts(self):
        if not self._mtdparts:
            pattern = b"mtdparts=mtdparts=alt11xx_sflash:"
            with open(self._file, "br") as fin:
                for chunk in self.each_chunk(fin, separator=b'\0'):
                    try:
                        chunk = chunk[chunk.index(pattern) + len(pattern):]
                        self.parse(chunk)
                        break
                    except ValueError:
                        pass

        return self._mtdparts


if __name__ == "__main__":
    uboot = Uboot(sys.argv[1])
    flashdump = Flashdump(uboot.mtdparts())
    print(flashdump)
    flashdump.split(sys.argv[2])