Python script for managing rigol dho800 partitions after dumping them

partitions.py

#!/usr/bin/env python

"""
This file is intended to be used after making a block-by-block copy of the micro sd card.
The partitions.txt file can be obtained by running `cat /proc/cmdline` on the oszi and then formatting it a bit.
"""

import os
import re

BACKUP_PATH = 'backups/backup.img'
PARTS_FOLDER = 'parts'
MULTIPLE_OF = 0x200 # 512 bytes
OFFSET_OFFSET = 0x2000
SKIP_NAMES = ['userdata']

# check if backup exists
if not os.path.exists(BACKUP_PATH):
    raise FileNotFoundError(f'Backup file not found: {BACKUP_PATH}')

# check if parts folder exists
if not os.path.exists(PARTS_FOLDER):
    os.mkdir(PARTS_FOLDER)

with open(BACKUP_PATH, 'rb') as img:
    idx = 1

    with open('partitions.txt', 'r') as f:
        for line in f:
            stripped = line.strip()

            if stripped.startswith('#'):
                continue

            if not stripped:
                continue

            # format: hex@hex(str)
            match = re.match(r'(0x[0-9a-fA-F]+|-)@(0x[0-9a-fA-F]+)\((.+)\)', stripped)
            if not match:
                raise ValueError(f'Invalid line: {stripped}')
            
            size = match.group(1)
            offset = match.group(2)
            name = match.group(3)

            if name in SKIP_NAMES:
                print(f'Skipping {name}')
                continue

            try:
                size = int(size, 16)
            except ValueError:
                if size != '-':
                    raise ValueError(f'Invalid size: {size}')
                size = None

            try:
                offset = int(offset, 16)
            except ValueError:
                raise ValueError(f'Invalid offset: {offset}')

            # split backup according to partitions.
            img.seek((offset + OFFSET_OFFSET) * MULTIPLE_OF)

            size = size * MULTIPLE_OF if size else None
            
            with open(os.path.join(PARTS_FOLDER, f"{idx:02d}_{name}.bin"), 'wb') as part:
                print(f'Writing partition {name} to {os.path.join(PARTS_FOLDER, name)} (pos: {img.tell()})')
                while True:
                    if size is not None:
                        if size <= 0:
                            break
                        chunk = img.read(min(size, 1024 * 1024))
                        size -= len(chunk)
                    else:
                        chunk = img.read(1024 * 1024)
                        if not chunk:
                            break
                    part.write(chunk)

            print(f'Partition {name} written to {os.path.join(PARTS_FOLDER, name)}')
            idx += 1