fix sdcard squashfs partitions for hassio

fix-hassio-sdcard

#!/usr/bin/env python3
"""
Script I with chatgpt wrote to help myself with hassio sd-card
squashfs images going broken as I decribed on
https://community.home-assistant.io/t/how-to-replace-broken-squashfs-images/738391/5

The idea is to specify sdcard and file image, verify that they are
identical in partitioning (besides last filesystem), and then copy some
partitions from image into the card.  Here is protocol of the use:

- first reviewed --dry-run:

    ❯ ~/bin/fix-hassio-sdcard --dry-run /dev/mmcblk0 haos_rpi3-64-12.3.img 1-6
    2024-06-11 16:30:54,060 - INFO - Running command: sudo dd if=haos_rpi3-64-12.3.img of=/dev/mmcblk0p1 bs=512 skip=2048 count=65536
    Dry run: sudo dd if=haos_rpi3-64-12.3.img of=/dev/mmcblk0p1 bs=512 skip=2048 count=65536
    2024-06-11 16:30:54,060 - INFO - Running command: sudo dd if=haos_rpi3-64-12.3.img of=/dev/mmcblk0p2 bs=512 skip=67584 count=49152
    Dry run: sudo dd if=haos_rpi3-64-12.3.img of=/dev/mmcblk0p2 bs=512 skip=67584 count=49152
    2024-06-11 16:30:54,060 - INFO - Running command: sudo dd if=haos_rpi3-64-12.3.img of=/dev/mmcblk0p3 bs=512 skip=116736 count=524288
    Dry run: sudo dd if=haos_rpi3-64-12.3.img of=/dev/mmcblk0p3 bs=512 skip=116736 count=524288
    2024-06-11 16:30:54,060 - INFO - Running command: sudo dd if=haos_rpi3-64-12.3.img of=/dev/mmcblk0p4 bs=512 skip=641024 count=49152
    Dry run: sudo dd if=haos_rpi3-64-12.3.img of=/dev/mmcblk0p4 bs=512 skip=641024 count=49152
    2024-06-11 16:30:54,060 - INFO - Running command: sudo dd if=haos_rpi3-64-12.3.img of=/dev/mmcblk0p5 bs=512 skip=690176 count=524288
    Dry run: sudo dd if=haos_rpi3-64-12.3.img of=/dev/mmcblk0p5 bs=512 skip=690176 count=524288
    2024-06-11 16:30:54,060 - INFO - Running command: sudo dd if=haos_rpi3-64-12.3.img of=/dev/mmcblk0p6 bs=512 skip=1214464 count=16384
    Dry run: sudo dd if=haos_rpi3-64-12.3.img of=/dev/mmcblk0p6 bs=512 skip=1214464 count=16384

- then applied it

    ❯ ~/bin/fix-hassio-sdcard  /dev/mmcblk0 haos_rpi3-64-12.3.img 1-6
    2024-06-11 16:31:28,022 - INFO - Running command: sudo dd if=haos_rpi3-64-12.3.img of=/dev/mmcblk0p1 bs=512 skip=2048 count=65536
    2024-06-11 16:31:32,799 - INFO - Running command: sudo dd if=haos_rpi3-64-12.3.img of=/dev/mmcblk0p2 bs=512 skip=67584 count=49152
    2024-06-11 16:31:36,469 - INFO - Running command: sudo dd if=haos_rpi3-64-12.3.img of=/dev/mmcblk0p3 bs=512 skip=116736 count=524288
    2024-06-11 16:32:14,562 - INFO - Running command: sudo dd if=haos_rpi3-64-12.3.img of=/dev/mmcblk0p4 bs=512 skip=641024 count=49152
    2024-06-11 16:32:18,212 - INFO - Running command: sudo dd if=haos_rpi3-64-12.3.img of=/dev/mmcblk0p5 bs=512 skip=690176 count=524288
    2024-06-11 16:32:57,107 - INFO - Running command: sudo dd if=haos_rpi3-64-12.3.img of=/dev/mmcblk0p6 bs=512 skip=1214464 count=16384

and hassio booted from sdcard again with my environment seems to be not impacted.

"""
import subprocess
import sys
import re
import argparse
import logging

def run_fdisk(device):
    result = subprocess.run(['sudo', 'fdisk', '-l', device], capture_output=True, text=True)
    if result.returncode != 0:
        logging.error(f"Error running fdisk on {device}: {result.stderr}")
        sys.exit(1)
    return result.stdout

def parse_fdisk_output(output):
    partitions = []
    lines = output.splitlines()
    sector_size = None
    for line in lines:
        if 'Sector size' in line:
            sector_size = int(re.search(r'(\d+) bytes', line).group(1))
        match = re.match(r'^(\/dev\/\S+|\S+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)', line)
        if match:
            partitions.append({
                'device': match.group(1),
                'start': int(match.group(2)),
                'end': int(match.group(3)),
                'sectors': int(match.group(4)),
                'size': match.group(5)
            })
    return partitions, sector_size

def check_partition_consistency(source_partitions, target_partitions):
    if len(source_partitions) != len(target_partitions):
        return False
    for i, (sp, tp) in enumerate(zip(source_partitions, target_partitions)):
        # Skip last one!
        if i == len(source_partitions)-1:
            logging.debug(f"Skipping partition {i}")
            break
        if sp['start'] != tp['start'] or sp['end'] != tp['end'] or sp['sectors'] != tp['sectors']:
            logging.error(f"Partition {i}: source={sp} target={tp}")
            return False
    return True

def copy_filesystem(image, partition, sector_size, dry_run):
    start = partition['start']
    sectors = partition['sectors']
    output_device = partition['device']
    dd_command = ['sudo', 'dd', f'if={image}', f'of={output_device}', f'bs={sector_size}', f'skip={start}', f'count={sectors}']
    logging.info(f"Running command: {' '.join(dd_command)}")
    if dry_run:
        print(f"Dry run: {' '.join(dd_command)}")
    else:
        result = subprocess.run(dd_command, capture_output=True, text=True)
        if result.returncode != 0:
            logging.error(f"Error running dd: {result.stderr}")
            sys.exit(1)

def main():
    parser = argparse.ArgumentParser(description='Copy filesystems from an image file to a target device.')
    parser.add_argument('target_device', help='Target device (e.g. /dev/mmcblk0)')
    parser.add_argument('source_image', help='Source image file (e.g. haos_rpi3-64-12.3.img)')
    parser.add_argument('range', help='Range of filesystem indexes to copy (e.g. 1-6)')
    parser.add_argument('--dry-run', action='store_true', help='Print the dd commands without executing them')
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

    target_device = args.target_device
    source_image = args.source_image
    range_indexes = args.range
    dry_run = args.dry_run

    try:
        start_index, end_index = map(int, range_indexes.split('-'))
    except ValueError:
        logging.error("Invalid range format. Please use the format start-end (e.g. 1-6).")
        sys.exit(1)

    target_fdisk_output = run_fdisk(target_device)
    source_fdisk_output = run_fdisk(source_image)

    target_partitions, target_sector_size = parse_fdisk_output(target_fdisk_output)
    source_partitions, source_sector_size = parse_fdisk_output(source_fdisk_output)

    if target_sector_size != source_sector_size:
        logging.error("Sector sizes do not match between target device and source image.")
        sys.exit(1)

    if not check_partition_consistency(source_partitions, target_partitions):
        logging.error("Partition tables do not match between target device and source image.")
        sys.exit(1)

    for i in range(start_index-1, end_index):
        if i >= len(source_partitions):
            logging.error(f"Invalid index: {i+1}. The source image has only {len(source_partitions)} partitions.")
            sys.exit(1)
        copy_filesystem(source_image, target_partitions[i], source_sector_size, dry_run)

if __name__ == '__main__':
    main()