ARK-OS/tools/pack_rpi4.py

#
# Copyright 2026 Aarav Ravindra Kharade
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#


#!/usr/bin/env python3
"""
ArkOS Raspberry Pi 4 Model B -- SD Card Image Builder
------------------------------------------------------
Generates a flashable SD card image (rpi4.img) containing:
  - Partition 1 (FAT32, 256MB): Boot partition
      kernel8.img, initramfs.img, bcm2711-rpi-4-b.dtb,
      config.txt, cmdline.txt
  - Partition 2 (ext4, 2GB): System partition (frameworks, libraries)
  - Partition 3 (ext4, 100MB): Vendor partition (DRM, keys, mirror)

The GPU firmware files (start4.elf, fixup4.dat, bootcode.bin) are NOT
included -- they are already present on every Raspberry Pi 4's EEPROM
and SD card from the factory.

Usage:
    python3 pack_rpi4.py --base-dir <arkos-dir>

Output:
    <arkos-dir>/finished/rpi4/rpi4.img
"""

import argparse
import os
import subprocess
import sys
import struct
import shutil


def run(cmd, check=True):
    """Execute a shell command and print it."""
    print(f"  $ {cmd}")
    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    if result.returncode != 0 and check:
        print(f"  ERROR: {result.stderr.strip()}")
        sys.exit(1)
    return result


def ensure_dir(path):
    os.makedirs(path, exist_ok=True)


def file_exists(path, name):
    """Check that a required file exists."""
    if not os.path.isfile(path):
        print(f"  ERROR: Missing required file: {path}")
        print(f"         ({name})")
        return False
    return True


def safe_copy(src, dst):
    """Safely copy src to dst, removing existing dst file if present."""
    if os.path.exists(dst):
        try:
            os.chmod(dst, 0o666)
            os.remove(dst)
        except Exception:
            pass
    shutil.copyfile(src, dst)
    try:
        os.chmod(dst, 0o666)
    except Exception:
        pass


def main():
    parser = argparse.ArgumentParser(description="ArkOS RPi4 SD Card Image Builder")
    parser.add_argument("--base-dir", required=True, help="Path to arkos/ directory")
    args = parser.parse_args()

    base = os.path.abspath(args.base_dir)
    staging = os.path.join(base, "out_staging", "rpi4")
    finished = os.path.join(base, "finished", "rpi4")
    boot_cfg = os.path.join(base, "boot", "rpi4")

    ensure_dir(staging)
    ensure_dir(finished)

    print("\n  RPi4 SD Card Image Builder")
    print("  " + "=" * 40)

    # ── Validate required files ──────────────────────────────────
    print("\n  Checking required files...")

    kernel_img = os.path.join(base, "kernel", "prebuilts", "Image")
    if not os.path.isfile(kernel_img):
        kernel_img = os.path.join(base, "kernel", "prebuilts", "arm64")

    initramfs = os.path.join(finished, "initramfs.img")
    # If initramfs is not in finished/rpi4/, check the main finished/ dir
    if not os.path.isfile(initramfs):
        initramfs = os.path.join(base, "finished", "initramfs.img")
    dtb = os.path.join(base, "kernel", "arch", "arm64", "boot", "dts",
                       "broadcom", "bcm2711-rpi-4-b.dtb")
    config_txt = os.path.join(boot_cfg, "config.txt")
    cmdline_txt = os.path.join(boot_cfg, "cmdline.txt")
    sys_img = os.path.join(finished, "sys.img")
    if not os.path.isfile(sys_img):
        sys_img = os.path.join(base, "finished", "sys.img")
        
    vend_img = os.path.join(finished, "vend.img")
    if not os.path.isfile(vend_img):
        vend_img = os.path.join(base, "finished", "vend.img")

    has_kernel = os.path.isfile(kernel_img)
    if not has_kernel:
        print(f"  WARNING: ARM64 kernel image missing at {kernel_img}")
        print(f"           Compile kernel with: make ARCH=arm64 CROSS_COMPILE=aarch64-linux-gnu- Image")
        print(f"           and place output at kernel/prebuilts/Image")

    has_initramfs = os.path.isfile(initramfs)
    if not has_initramfs:
        print(f"  WARNING: initramfs missing at {initramfs}")

    ok = file_exists(config_txt, "RPi4 config.txt") and file_exists(cmdline_txt, "RPi4 cmdline.txt")

    # DTB is optional if user compiles kernel separately
    has_dtb = os.path.isfile(dtb)
    if not has_dtb:
        # Try prebuilts directory
        dtb_alt = os.path.join(base, "kernel", "prebuilts", "bcm2711-rpi-4-b.dtb")
        if os.path.isfile(dtb_alt):
            dtb = dtb_alt
            has_dtb = True
        else:
            print(f"  WARNING: Device tree blob not found at {dtb}")
            print(f"           or {dtb_alt}")
            print(f"           Compile DTB with: make ARCH=arm64 CROSS_COMPILE=aarch64-linux-gnu- dtbs")

    if not ok:
        print("\n  Build failed: missing required boot configuration files.")
        sys.exit(1)

    # ── Stage GPU Firmware & Boot Files ───────────────────────────
    fw_dir = os.path.join(boot_cfg, "firmware")
    ensure_dir(fw_dir)

    start4 = os.path.join(fw_dir, "start4.elf")
    fixup4 = os.path.join(fw_dir, "fixup4.dat")

    # If GPU firmware files are missing, download official release from raspberrypi/firmware
    if not os.path.isfile(start4) or not os.path.isfile(fixup4):
        print(f"  Fetching RPi4 GPU firmware (start4.elf, fixup4.dat)...")
        try:
            import urllib.request
            fw_repo = "https://raw.githubusercontent.com/raspberrypi/firmware/master/boot"
            if not os.path.isfile(start4):
                urllib.request.urlretrieve(f"{fw_repo}/start4.elf", start4)
                print(f"  Downloaded start4.elf")
            if not os.path.isfile(fixup4):
                urllib.request.urlretrieve(f"{fw_repo}/fixup4.dat", fixup4)
                print(f"  Downloaded fixup4.dat")
        except Exception as e:
            print(f"  WARNING: Could not download GPU firmware files automatically ({e}).")
            print(f"           Place start4.elf and fixup4.dat in boot/rpi4/firmware/")

    # ── Stage boot partition files ───────────────────────────────
    print("\n  Staging boot partition files...")
    boot_stage = os.path.join(staging, "boot")
    shutil.rmtree(boot_stage, ignore_errors=True)
    ensure_dir(boot_stage)

    # Copy GPU firmware if present
    if os.path.isfile(start4):
        safe_copy(start4, os.path.join(boot_stage, "start4.elf"))
        print(f"  Copied start4.elf")
    if os.path.isfile(fixup4):
        safe_copy(fixup4, os.path.join(boot_stage, "fixup4.dat"))
        print(f"  Copied fixup4.dat")

    # Copy kernel as kernel8.img (RPi4 convention for 64-bit)
    if has_kernel:
        safe_copy(kernel_img, os.path.join(boot_stage, "kernel8.img"))
        print(f"  Copied kernel Image -> kernel8.img")

    # Copy initramfs
    if has_initramfs:
        safe_copy(initramfs, os.path.join(boot_stage, "initramfs.img"))
        print(f"  Copied initramfs.img")

    # Copy config.txt and cmdline.txt
    safe_copy(config_txt, os.path.join(boot_stage, "config.txt"))
    safe_copy(cmdline_txt, os.path.join(boot_stage, "cmdline.txt"))
    print(f"  Copied config.txt, cmdline.txt")

    # Copy DTB if available
    if has_dtb:
        safe_copy(dtb, os.path.join(boot_stage, "bcm2711-rpi-4-b.dtb"))
        print(f"  Copied bcm2711-rpi-4-b.dtb")

    # Copy overlays directory if available
    dtbo_dir = os.path.join(base, "kernel", "prebuilts", "rpi4_dtbofiles")
    if os.path.isdir(dtbo_dir):
        boot_overlays = os.path.join(boot_stage, "overlays")
        ensure_dir(boot_overlays)
        for f in os.listdir(dtbo_dir):
            if f.endswith(".dtbo") or f.endswith(".dtb"):
                safe_copy(os.path.join(dtbo_dir, f), os.path.join(boot_overlays, f))
        print(f"  Copied overlays/ to boot partition")

    # ── Create boot partition image (FAT32, 256MB) ──────────────
    print("\n  Creating FAT32 boot partition (256MB)...")
    boot_img = os.path.join(staging, "boot.img")
    run(f"dd if=/dev/zero of={boot_img} bs=1M count=256 2>/dev/null")
    run(f"mkfs.vfat -F 32 -n ARKOS_BOOT {boot_img} 2>/dev/null")

    # Copy files into the FAT32 image using mcopy (mtools)
    for fname in os.listdir(boot_stage):
        src = os.path.join(boot_stage, fname)
        if os.path.isdir(src):
            run(f"mcopy -s -i {boot_img} {src} ::/{fname}")
        else:
            run(f"mcopy -i {boot_img} {src} ::/{fname}")
    print(f"  Boot partition ready ({os.path.getsize(boot_img)} bytes)")

    # ── Assemble final SD card image ─────────────────────────────
    print("\n  Assembling SD card image...")

    # Calculate sizes
    boot_size_mb = 256
    sys_size_mb = 2048
    vend_size_mb = 100
    total_mb = boot_size_mb + sys_size_mb + vend_size_mb + 6  # +6MB for GPT headers

    rpi4_img = os.path.join(finished, "rpi4.img")

    # Create empty image
    run(f"dd if=/dev/zero of={rpi4_img} bs=1M count={total_mb} 2>/dev/null")

    # Create GPT partition table with 3 partitions
    # Partition 1: FAT32 boot (256MB)
    # Partition 2: ext4 system (2GB)
    # Partition 3: ext4 vendor (100MB)
    vend_end = boot_size_mb + sys_size_mb + vend_size_mb + 4
    run(f"parted -s {rpi4_img} mklabel gpt")
    run(f"parted -s {rpi4_img} mkpart boot fat32 4MiB {boot_size_mb + 4}MiB")
    run(f"parted -s {rpi4_img} mkpart system ext4 {boot_size_mb + 4}MiB {boot_size_mb + sys_size_mb + 4}MiB")
    run(f"parted -s {rpi4_img} mkpart vendor ext4 {boot_size_mb + sys_size_mb + 4}MiB {vend_end}MiB")
    run(f"parted -s {rpi4_img} set 1 boot on")

    # Write boot partition at offset 4MiB
    boot_offset = 4 * 1024 * 1024
    run(f"dd if={boot_img} of={rpi4_img} bs=1M seek=4 conv=notrunc 2>/dev/null")

    # Write system partition
    if os.path.isfile(sys_img):
        sys_offset = (boot_size_mb + 4)
        run(f"dd if={sys_img} of={rpi4_img} bs=1M seek={sys_offset} conv=notrunc 2>/dev/null")
        print(f"  Wrote system partition at offset {sys_offset}MB")
    else:
        print(f"  WARNING: sys.img not found, system partition will be empty")

    # Write vendor partition
    if os.path.isfile(vend_img):
        vend_offset = (boot_size_mb + sys_size_mb + 4)
        run(f"dd if={vend_img} of={rpi4_img} bs=1M seek={vend_offset} conv=notrunc 2>/dev/null")
        print(f"  Wrote vendor partition at offset {vend_offset}MB")
    else:
        print(f"  WARNING: vend.img not found, vendor partition will be empty")

    # ── Summary ──────────────────────────────────────────────────
    img_size = os.path.getsize(rpi4_img)
    print(f"\n  SD card image built successfully!")
    print(f"  Output: {rpi4_img}")
    print(f"  Size:   {img_size:,} bytes ({img_size // (1024*1024)} MB)")
    print(f"\n  Flash to SD card:")
    print(f"    sudo dd if={rpi4_img} of=/dev/sdX bs=4M status=progress")
    print(f"    sync")
    print()


if __name__ == "__main__":
    main()