setup_btrfs_vols.py

#!/usr/bin/env python3

"""
Automates creation of a baseline Btrfs layout on a given device/partition.
- Assumes the partition is already wiped or otherwise ready.
- Creates a Btrfs filesystem with a label.
- Creates subvolumes: @, @boot, @swap, @snapshots, @var, @home, @tmp
- Mounts @ at /mnt, then mounts the others to their respective subdirectories.
- NOTE: You must run this script with sudo (or as root).

Usage:
    python setup_btrfs.py --device /dev/nvme0n1p1 --label MyBTRFSLabel

After completion:
    - /mnt will contain @ (root subvol)
    - /mnt/boot -> @boot
    - /mnt/swap -> @swap
    - /mnt/snapshots -> @snapshots
    - /mnt/var -> @var
    - /mnt/home -> @home
    - /mnt/tmp -> @tmp

You can then proceed to install Ubuntu into /mnt (chroot) or adjust fstab, etc.
"""

import argparse
import subprocess
import sys
import os


def run_cmd(cmd, check=True):
    """
    Helper to run a command using subprocess.
    If check=True, raises an exception on nonzero return.
    """
    print(f"[CMD] {' '.join(cmd)}")
    subprocess.run(cmd, check=check)


def get_uuid(device):
    """
    Retrieve the UUID of a device using blkid.
    """

    result = subprocess.run(
        ["blkid", "-s", "UUID", "-o", "value", device],
        capture_output=True,
        text=True,
        check=True,
    )
    return result.stdout.strip()


def create_swap_file(swap_mount_dir, swap_size):
    """
    Create and configure a swap file at 'swap_mount_dir/swapfile' of size 'swap_size'
    - Disables CoW via chattr + C
    - Creae the swap file
    - executes mkswap, chmod
    """

    swap_file_path = os.path.join(swap_mount_dir, "swapfile")

    # 1. Touch an empty file, disable COW immediately
    run_cmd(["touch", swap_file_path])
    run_cmd(["chattr", "+C", swap_file_path])

    # 2. Populate the file. Using fallocate is faster than dd.
    run_cmd(["fallocate", "-l", swap_size, swap_file_path])

    # 3. Set correct permnission, then run mkswap
    run_cmd(["chmod", "600", swap_file_path])
    run_cmd(["mkswap", swap_file_path])


def append_fstab_entries(mount_path, device_uuid):
    """
    Append btrfs subvolume entries and swapfile entry to /mnt/etc/fstab
    """

    etc_dir = os.path.join(mount_path, "etc")
    os.makedirs(etc_dir, exist_ok=True)

    fstab_path = os.path.join(etc_dir, "fstab")

    # We assume user wants stadnard default mounts
    # We may want to add other options like

    fstab_entries = f"""
    # BTRFS subvolume mounts
    UUID={device_uuid}      /               btrfs defaults,subvol=@                 0 0
    UUID={device_uuid}      /boot           btrfs defaults,subvol=@boot             0 0
    UUID={device_uuid}      /swap           btrfs defaults,subvol=@swap             0 0
    UUID={device_uuid}      /snapshots      btrfs defaults,subvol=@snapshots        0 0
    UUID={device_uuid}      /var            btrfs defaults,subvol=@var              0 0
    UUID={device_uuid}      /home           btrfs defaults,subvol=@home             0 0
    UUID={device_uuid}      /tmp            btrfs defaults,subvol=@tmp              0 0

    # Swap file
    /swap/swapfile          none            swap    sw                              0 0

    """

    print("Appending the following to /mnt/etc/fstab:\n")
    print(fstab_entries)

    with open(fstab_path, "a") as f:
        f.write(fstab_entries)


def main():
    parser = argparse.ArgumentParser(
        description="Automate Btrfs creation and subvolume setup on Ubuntu."
    )
    parser.add_argument(
        "--device", required=True, help="Path to the partition (e.g. /dev/nvme0n1p1)."
    )
    parser.add_argument(
        "--label", default="UBUNTU_BTRFS", help="Label to set on the Btrfs filesystem."
    )
    parser.add_argument(
        "--swap-size",
        required=True,
        help="Size of the swap file to create (e.g. 4GB or 8192MB etc.).",
    )
    args = parser.parse_args()

    device = args.device
    label = args.label
    swap_size = args.swap_size

    # 1. Create the Btrfs filesystem on the device.
    #    (User is assumed to have wiped the disk/partition beforehand.)
    print(f"Formatting {device} as Btrfs with label '{label}'...")
    run_cmd(["mkfs.btrfs", "-f", "-L", label, device])

    # 2. Create a temporary mount point for the top-level (ID 5).
    top_level_mount = "/mnt/btrfs_top_level"
    os.makedirs(top_level_mount, exist_ok=True)

    print("Mounting top-level subvolume (ID 5)...")
    run_cmd(["mount", "-o", "subvolid=5", device, top_level_mount])

    # 3. Create the desired subvolumes.
    subvolumes = ["@", "@boot", "@swap", "@snapshots", "@var", "@home", "@tmp"]
    print("Creating subvolumes...")
    for sv in subvolumes:
        run_cmd(["btrfs", "subvolume", "create", os.path.join(top_level_mount, sv)])

    # 4. Unmount the top-level.
    print("Unmounting top-level subvolume...")
    run_cmd(["umount", top_level_mount])
    os.rmdir(top_level_mount)

    # 5. Mount the root subvolume @ at /mnt.
    os.makedirs("/mnt", exist_ok=True)
    print("Mounting @ (root subvolume) at /mnt...")
    run_cmd(["mount", "-o", "subvol=@", device, "/mnt"])

    # 6. Create mount points for the other subvolumes within /mnt.
    #    Then mount each subvolume in the correct place.
    mount_map = {
        "@boot": "/mnt/boot",
        "@swap": "/mnt/swap",
        "@snapshots": "/mnt/snapshots",
        "@var": "/mnt/var",
        "@home": "/mnt/home",
        "@tmp": "/mnt/tmp",
    }
    print("Creating directories for additional subvolumes...")
    for subvol, mntp in mount_map.items():
        os.makedirs(mntp, exist_ok=True)

    print("Mounting additional subvolumes...")
    for subvol, mntp in mount_map.items():
        run_cmd(["mount", "-o", f"subvol={subvol}", device, mntp])

    # 7. Create the swap file in the @swap subvolume & configure as needed
    print(f"Create a swap file of size {swap_size} at /mnt/swap/swapfile...")
    create_swap_file("/mnt/swap", swap_size)

    # 8. Get device UUID and append the approp. lines to /mnt/etc/fstab
    print("Retrieving UUID of the device & Updating fstab entries...")
    device_uuid = get_uuid(device)
    append_fstab_entries("/mnt", device_uuid)

    # 9. Done: We have a baseline Btrfs layout at /mnt.
    print()
    print("Btrfs setup complete!")
    print("Layout is now mounted at /mnt.")
    print(
        "You can now proceed to install Ubuntu into /mnt (chroot) or adjust fstab, etc."
    )
    print()


if __name__ == "__main__":
    if os.geteuid() != 0:
        print("This script must be run as root (sudo). Exiting.")
        sys.exit(1)
    main()