create-linux-swapfile/src/swap_forge/core.py

import re
import subprocess
from pathlib import Path

FSTAB_FILE = Path("/etc/fstab")
BTRFS_SWAP_DIR = Path("/var/swap")
BTRFS_SWAP_PATH = BTRFS_SWAP_DIR / "swapfile"
DEFAULT_SWAP_PATH = Path("/swapfile")

SIZE_PATTERN = re.compile(r"^(\d+)([gGmM]?[bB]?)?$")
MIB_PER_GIB = 1024

REQUIRED_BINARIES = ("findmnt", "swapoff", "chmod", "mkswap", "swapon")
BTRFS_BINARIES = ("chattr", "btrfs", "dd")
NON_BTRFS_BINARIES = ("fallocate",)
EXIT_MISSING_BINARY = 127
EXIT_COMMAND_FAILED = 1


class MissingBinaryError(RuntimeError):
    """A required external command is not installed.

    Args:
        binary: Name of the missing executable.
    """

    def __init__(self, binary):
        super().__init__(f"required command '{binary}' is not installed")
        self.binary = binary


def run(cmd, check=True, capture=False):
    """Run a command and return the completed process.

    Args:
        cmd: Command argument list.
        check: Raise CalledProcessError when the command fails.
        capture: Capture stdout and stderr as text.
    """
    try:
        if capture:
            return subprocess.run(cmd, check=check, capture_output=True, text=True)
        return subprocess.run(cmd, check=check)
    except FileNotFoundError as error:
        raise MissingBinaryError(cmd[0]) from error


def detect_root_fs():
    """Return the filesystem type mounted at /."""
    return run(["findmnt", "-no", "FSTYPE", "/"], capture=True).stdout.strip()


def parse_size_to_mib(size):
    """Convert a size such as '64G' or '2048M' into MiB.

    Args:
        size: Size string with an optional G/M unit, defaulting to MiB.
    """
    match = SIZE_PATTERN.match(size.strip())
    if not match:
        raise ValueError(f"Invalid size format: {size}")

    number, unit = match.groups()
    unit = (unit or "M").lower()

    if unit in ("g", "gb"):
        return int(number) * MIB_PER_GIB
    if unit in ("m", "mb"):
        return int(number)
    raise ValueError(f"Unsupported unit: {unit}")


def get_swap_size(path):
    """Return the size of the swapfile in MiB, or 0 when it does not exist.

    Args:
        path: Path of the swapfile.
    """
    if not path.exists():
        return 0
    return path.stat().st_size // (1024 * 1024)


def swap_fstab_line(path):
    """Return the fstab line registering the given swapfile."""
    return f"{path} none swap sw 0 0"


def remove_swap(path, fstab=FSTAB_FILE):
    """Disable the swapfile and drop it plus its fstab entry.

    Args:
        path: Path of the swapfile.
        fstab: Path of the fstab file to clean up.
    """
    run(["swapoff", str(path)], check=False)
    if path.exists():
        path.unlink()

    if not fstab.exists():
        return

    kept = [
        line
        for line in fstab.read_text().splitlines()
        if not re.search(rf"^{re.escape(str(path))}\s+none\s+swap", line)
    ]
    fstab.write_text("\n".join(kept) + "\n")


def create_swap(path, size_mib, fs, fstab=FSTAB_FILE):
    """Create, activate and register a swapfile.

    Args:
        path: Path of the swapfile.
        size_mib: Desired size in MiB.
        fs: Filesystem type of the root mount.
        fstab: Path of the fstab file to append to.
    """
    if fs == "btrfs":
        print(f"Creating {size_mib} MiB swapfile on btrfs at {path}")
        path.parent.mkdir(parents=True, exist_ok=True)
        run(["chattr", "+C", str(path.parent)], check=False)
        run(
            [
                "btrfs",
                "property",
                "set",
                "-ts",
                str(path.parent),
                "compression",
                "none",
            ],
            check=False,
        )
        run(
            [
                "dd",
                "if=/dev/zero",
                f"of={path}",
                "bs=1M",
                f"count={size_mib}",
                "status=progress",
            ]
        )
    else:
        print(f"Creating {size_mib} MiB swapfile on {fs} at {path}")
        run(["fallocate", "-l", f"{size_mib}M", str(path)])

    run(["chmod", "600", str(path)])
    run(["mkswap", str(path)])
    run(["swapon", str(path)])

    with fstab.open("a") as handle:
        handle.write(swap_fstab_line(path) + "\n")


def swap_path_for(fs):
    """Return the swapfile path matching the given root filesystem type."""
    return BTRFS_SWAP_PATH if fs == "btrfs" else DEFAULT_SWAP_PATH