feat: add cachy-chroot script (#62)

Initial release of cachy-chroot script. This script simplifies the
process of chrooting by prompt-based selection of the block devices and
the mount points. It also handles BTRFS subvolumes automatically.

Signed-off-by: SoulHarsh007 <[email protected]>

usr/bin/cachy-chroot

@@ -0,0 +1,370 @@
+#!/usr/bin/env python3
+
+"""
+CachyOS Chroot Helper
+
+This script is a helper for chrooting into a CachyOS installation. It will help you
+mount the root partition and any additional partitions you want to mount before
+chrooting into the system.
+"""
+
+__version__ = "1.0.0"
+__author__ = "SoulHarsh007"
+__license__ = "BSD-3-Clause"
+__maintainer__ = "SoulHarsh007"
+__email__ = "[email protected]"
+
+import argparse
+import json
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+from typing import Dict, List
+
+parser = argparse.ArgumentParser(description="Chroot helper for CachyOS")
+
+parser.add_argument(
+    "--skip-root-check",
+    help="Skip root check",
+    action="store_true",
+)
+
+args = parser.parse_args()
+
+if not args.skip_root_check and os.geteuid() != 0:
+    print(
+        "This script must be run as root to work properly. Use --skip-root-check to skip this check."
+    )
+    sys.exit(1)
+
+lsblk = shutil.which("lsblk")
+mount = shutil.which("mount")
+umount = shutil.which("umount")
+btrfs = shutil.which("btrfs")
+arch_chroot = shutil.which("arch-chroot")
+
+if not lsblk:
+    print("lsblk not found, please install it to continue. (e.g. pacman -S util-linux)")
+    sys.exit(1)
+
+if not mount:
+    print("mount not found, please install it to continue. (e.g. pacman -S util-linux)")
+    sys.exit(1)
+
+if not umount:
+    print(
+        "umount not found, please install it to continue. (e.g. pacman -S util-linux)"
+    )
+    sys.exit(1)
+
+if not btrfs:
+    print(
+        "btrfs not found, please install it to continue. (e.g. pacman -S btrfs-progs)"
+    )
+    sys.exit(1)
+
+if not arch_chroot:
+    print(
+        "arch-chroot not found, please install it to continue. (e.g. pacman -S arch-install-scripts)"
+    )
+    sys.exit(1)
+
+data, err = subprocess.Popen(
+    [lsblk, "-f", "-o", "NAME,FSTYPE,TYPE,UUID", "-p", "-a", "-J"],
+    stdout=subprocess.PIPE,
+).communicate()
+
+if err:
+    print("Failed to list block devices")
+    print(err)
+    sys.exit(1)
+
+
+class BlockDevice:
+    def __init__(self, name: str, fstype: str, uuid: str):
+        self.name = name
+        self.fstype = fstype
+        self.uuid = uuid
+
+    def __str__(self):
+        return f"{self.name} ({self.fstype}) with UUID {self.uuid}"
+
+    def __repr__(self):
+        return self.__str__()
+
+
+class BTRFSSubvolume(BlockDevice):
+    def __init__(self, name: str, fstype: str, uuid: str, path: str, subvol_id: int):
+        super().__init__(name, fstype, uuid)
+        self.path = path
+        self.subvol_id = subvol_id
+
+    def __str__(self):
+        return f"{self.name} ({self.fstype}) with UUID {self.uuid} at {self.path} (subvol_id: {self.subvol_id})"
+
+
+data = json.loads(data)
+
+valid_devices: List[BlockDevice] = []
+
+for device in data["blockdevices"]:
+    if (
+        device["type"] == "disk"
+        and "children" in device
+        and len(device["children"]) > 0
+    ):
+        print(f"Found disk {device['name']} with partitions:")
+        for partition in device["children"]:
+            if partition["fstype"] != "swap":
+                block_device = BlockDevice(
+                    partition["name"], partition["fstype"], partition["uuid"]
+                )
+                valid_devices.append(block_device)
+                print(f"  {block_device}")
+            else:
+                print(f"  {partition['name']} (swap) (ignored)")
+
+num_devices = len(valid_devices)
+
+if num_devices == 0:
+    print("No valid block devices found")
+    sys.exit(1)
+
+
+def get_user_choice(
+    partition_name: str,
+    num_devices: int,
+    valid_devices: List[BlockDevice] | List[BTRFSSubvolume] = valid_devices,
+):
+    for index, device in enumerate(valid_devices):
+        print(f"  {index + 1}. {device}")
+    user_choice = input(f"Select your {partition_name} partition [1-{num_devices}]: ")
+    try:
+        user_choice = int(user_choice)
+        if user_choice < 1 or user_choice > num_devices:
+            raise ValueError
+    except ValueError:
+        print(
+            f"Invalid choice, please enter a number between 1 and {num_devices}. Try again."
+        )
+        return get_user_choice(partition_name, num_devices)
+    return valid_devices[user_choice - 1]
+
+
+def mount_block_device(
+    block_device: BlockDevice, mount_point: str, options: List[str] = None
+):
+    if options is None:
+        options = []
+    print(f"Mounting {block_device} to {mount_point} with options {options}")
+    os.makedirs(mount_point, exist_ok=True)
+    process = subprocess.run([mount, block_device.name, mount_point, *options])
+    process.check_returncode()
+
+
+def unmount_block_device(mount_point: str):
+    print(f"Unmounting {mount_point}")
+    process = subprocess.run([umount, mount_point])
+    process.check_returncode()
+
+
+def scan_btrfs_subvolumes(btrfs_partition: BlockDevice):
+    new_subvolumes: List[BTRFSSubvolume] = []
+    print(f"Scanning BTRFS subvolumes on {btrfs_partition}")
+    with tempfile.TemporaryDirectory(
+        f"cachy-chroot-temp-mount-{btrfs_partition.uuid}", dir=tempfile.gettempdir()
+    ) as temp_directory:
+        mount_block_device(btrfs_partition, temp_directory)
+        subvolumes = (
+            subprocess.run(
+                [btrfs, "subvolume", "list", temp_directory],
+                stdout=subprocess.PIPE,
+            )
+            .stdout.decode("utf-8")
+            .splitlines()
+        )
+        for subvolume in subvolumes:
+            subvol_id, path = subvolume.split(" path ")
+            subvol_id = int(subvol_id.split(" ")[1])
+            path = path.strip()
+            btrfs_subvolume = BTRFSSubvolume(
+                btrfs_partition.name,
+                btrfs_partition.fstype,
+                btrfs_partition.uuid,
+                path,
+                subvol_id,
+            )
+            print(f"  {btrfs_subvolume}")
+            new_subvolumes.append(btrfs_subvolume)
+        unmount_block_device(temp_directory)
+    return new_subvolumes
+
+
+root_partition: BlockDevice | BTRFSSubvolume = get_user_choice("root", num_devices)
+
+print(f"Selected root partition: {root_partition}")
+
+discovered_subvolumes: Dict[str, List[BTRFSSubvolume]] = {}
+
+
+def get_user_btrfs_subvolume_choice(
+    partition_name: str, num_subvolumes: int, subvolumes: List[BTRFSSubvolume]
+):
+    for index, subvolume in enumerate(subvolumes):
+        print(f"  {index + 1}. {subvolume}")
+    print(
+        f"  {num_subvolumes + 1}. Use {partition_name} partition as-is (subvol_id 5 @ /)"
+    )
+    user_choice = input(
+        f"Select your {partition_name} subvolume [1-{num_subvolumes + 1}]: "
+    )
+    try:
+        user_choice = int(user_choice)
+        if user_choice < 1 or user_choice > num_subvolumes + 1:
+            raise ValueError
+    except ValueError:
+        print(
+            f"Invalid choice, please enter a number between 1 and {num_subvolumes}. Try again."
+        )
+        return get_user_btrfs_subvolume_choice(
+            partition_name, num_subvolumes, subvolumes
+        )
+    if user_choice == num_subvolumes + 1:
+        root_partition = BTRFSSubvolume(
+            subvolumes[0].name, "btrfs", subvolumes[0].uuid, "@", 5
+        )
+        return root_partition
+    return subvolumes[user_choice - 1]
+
+
+if root_partition.fstype == "btrfs":
+    new_subvolumes: List[BTRFSSubvolume] = []
+    print(
+        "Detected BTRFS filesystem on root partition, mounting to temporary location to list subvolumes"
+    )
+    new_subvolumes = scan_btrfs_subvolumes(root_partition)
+    discovered_subvolumes[root_partition.name] = new_subvolumes
+    print(f"Found {len(new_subvolumes)} subvolumes on {root_partition}")
+    if len(new_subvolumes) > 0:
+        root_partition = get_user_btrfs_subvolume_choice(
+            "root", len(new_subvolumes), new_subvolumes
+        )
+    else:
+        print("No subvolumes found, using root partition as-is")
+        root_partition = BTRFSSubvolume(
+            root_partition.name, root_partition.fstype, root_partition.uuid, "@", 5
+        )
+
+temp_mount_dir = tempfile.mkdtemp(
+    f"cachy-chroot-working-mount-{root_partition.uuid}", dir=tempfile.gettempdir()
+)
+
+print(f"Final root partition: {root_partition}")
+
+mounted_partitions: List[BlockDevice | BTRFSSubvolume] = []
+
+if isinstance(root_partition, BTRFSSubvolume):
+    mount_block_device(
+        root_partition, temp_mount_dir, ["-o", f"subvolid={root_partition.subvol_id}"]
+    )
+elif root_partition.fstype == "btrfs":
+    mount_block_device(root_partition, temp_mount_dir, ["-o", "subvolid=5"])
+else:
+    mount_block_device(root_partition, temp_mount_dir)
+
+mounted_partitions.append(root_partition)
+
+
+# TODO: Parse /etc/fstab and mount block devices
+def parse_fstab():
+    if not os.path.exists(os.path.join(temp_mount_dir, "etc/fstab")):
+        print("warning: No /etc/fstab found, skipping auto-mounting")
+        print(
+            "Either the system is not installed or the root partition is not mounted properly. Proceed with caution."
+        )
+        return
+    print("Parsing /etc/fstab for auto-mounting")
+
+
+def mount_additional_block_devices():
+    print("What partition do you want to mount? (Example: /boot, /home, /var, etc.): ")
+    user_selection = input("Enter the partition path: ")
+    user_selection = user_selection.strip().lower()
+    if user_selection == "skip":
+        print("Skipping additional block device mounting...")
+        return
+    if not user_selection or not user_selection.startswith("/"):
+        print("Invalid partition path, please enter a valid path.")
+        print("Example: /boot, /home, /var, etc.")
+        print("Don't want to mount additional block devices? Enter 'skip' to skip.")
+        return mount_additional_block_devices()
+    block_device = get_user_choice("additional block device", num_devices)
+    if block_device in mounted_partitions:
+        print(f"{block_device} is already mounted, skipping...")
+        return
+    mount_point = os.path.join(temp_mount_dir, user_selection)
+    if block_device.fstype == "btrfs":
+        subvolumes: List[BTRFSSubvolume] = []
+        if block_device.name in discovered_subvolumes:
+            subvolumes = discovered_subvolumes[block_device.name]
+        else:
+            subvolumes = scan_btrfs_subvolumes(block_device)
+            discovered_subvolumes[block_device.name] = subvolumes
+        print(f"Found {len(subvolumes)} subvolumes on {block_device}")
+        if len(subvolumes) > 0:
+            block_device = get_user_btrfs_subvolume_choice(
+                "additional block device", len(subvolumes), subvolumes
+            )
+        else:
+            print("No subvolumes found, using block device as-is")
+            block_device = BTRFSSubvolume(
+                block_device.name, block_device.fstype, block_device.uuid, "@", 5
+            )
+    if isinstance(block_device, BTRFSSubvolume):
+        mount_block_device(
+            block_device, mount_point, ["-o", f"subvolid={block_device.subvol_id}"]
+        )
+    elif block_device.fstype == "btrfs":
+        mount_block_device(block_device, mount_point, ["-o", "subvolid=5"])
+    else:
+        mount_block_device(
+            block_device,
+            mount_point,
+        )
+    mounted_partitions.append(block_device)
+
+
+def mount_additional_block_devices_prompt():
+    print("Do you want to mount additional block devices? ")
+    user_response = input("Enter 'yes' to continue, or 'no' to skip: ")
+    user_response = user_response.strip().lower()
+    if user_response == "yes" or user_response == "y":
+        mount_additional_block_devices()
+        return mount_additional_block_devices_prompt()
+    elif user_response == "no" or user_response == "n":
+        print("Not mounting additional block devices...")
+    else:
+        print("Invalid response, please enter 'yes' or 'no'")
+        return mount_additional_block_devices_prompt()
+
+
+mount_additional_block_devices_prompt()
+print("Chrooting into the system...")
+print(
+    "Remember to exit the chroot environment when you're done with 'exit' or 'Ctrl+D'"
+)
+print("Happy chrooting! :)")
+process = subprocess.run([arch_chroot, temp_mount_dir])
+try:
+    process.check_returncode()
+except subprocess.CalledProcessError:
+    print("chroot process exited with non-zero exit code :(")
+    print("Please check the logs above for more information.")
+    sys.exit(1)
+finally:
+    print("Unmounting block devices...")
+    for mounted_partition in reversed(mounted_partitions):
+        unmount_block_device(mounted_partition.name)
+    os.rmdir(temp_mount_dir)