diff --git a/main.py b/main.py new file mode 100644 index 0000000..39771c5 --- /dev/null +++ b/main.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python3 +import argparse +import os +import re +import shutil +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") + + +def run(cmd, check=True, capture=False): + """Helper to run shell commands.""" + kwargs = {"check": check} + if capture: + kwargs["stdout"] = subprocess.PIPE + kwargs["stderr"] = subprocess.PIPE + kwargs["text"] = True + return subprocess.run(cmd, **kwargs) + + +def detect_root_fs() -> str: + result = run(["findmnt", "-no", "FSTYPE", "/"], capture=True) + return result.stdout.strip() + + +def parse_size_to_mib(size: str) -> int: + """Convert '64G' or '2048M' into MiB integer.""" + match = re.match(r"^(\d+)([gGmM]?[bB]?)?$", size.strip()) + if not match: + raise ValueError(f"Invalid size format: {size}") + num, unit = match.groups() + num = int(num) + unit = (unit or "M").lower() + + if unit in ("g", "gb"): + return num * 1024 + elif unit in ("m", "mb"): + return num + else: + raise ValueError(f"Unsupported unit: {unit}") + + +def get_swap_size(path: Path) -> int: + """Return swapfile size in MiB if it exists, else 0.""" + if not path.exists(): + return 0 + return path.stat().st_size // (1024 * 1024) + + +def remove_swap(path: Path): + """Disable and remove swapfile and its fstab entry.""" + run(["swapoff", str(path)], check=False) + if path.exists(): + path.unlink() + # remove old fstab line + if FSTAB_FILE.exists(): + text = FSTAB_FILE.read_text().splitlines() + new_lines = [ + line for line in text + if not re.search(rf"^{re.escape(str(path))}\s+none\s+swap", line) + ] + FSTAB_FILE.write_text("\n".join(new_lines) + "\n") + + +def create_swap(path: Path, size_mib: int, fs: str): + 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_FILE.open("a") as f: + f.write(f"{path} none swap sw 0 0\n") + + +def main(): + parser = argparse.ArgumentParser(description="SwapForge Python edition") + parser.add_argument("size", help="Swapfile size (e.g. 2048M or 64G)") + args = parser.parse_args() + + fs = detect_root_fs() + path = BTRFS_SWAP_PATH if fs == "btrfs" else DEFAULT_SWAP_PATH + new_size = parse_size_to_mib(args.size) + + old_size = get_swap_size(path) + if old_size == new_size and old_size > 0: + print(f"Swapfile {path} already exists with correct size ({old_size} MiB). Skipping.") + return + + if old_size > 0 and old_size != new_size: + print(f"Existing swapfile {path} has size {old_size} MiB, expected {new_size} MiB. Recreating...") + remove_swap(path) + + create_swap(path, new_size, fs) + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/main.sh b/main.sh deleted file mode 100644 index 27cbc6c..0000000 --- a/main.sh +++ /dev/null @@ -1,14 +0,0 @@ -#!/bin/bash -FSTAB_SWAP_ENTRY="/swapfile none swap defaults 0 0" -SWAP_FILE="/swapfile" -FSTAB_FILE="/etc/fstab" -if grep -q "$FSTAB_SWAP_ENTRY" "$FSTAB_FILE"; then - echo "Skipping creation of swap partion because entry allready exists in \"$FSTAB_FILE\"!" -else - echo "Creating swap partition..." && - sudo fallocate -l "$1" "$SWAP_FILE" && - sudo chmod 600 "$SWAP_FILE" && - sudo mkswap "$SWAP_FILE" && - sudo swapon "$SWAP_FILE" && - sudo sh -c "echo \"$FSTAB_SWAP_ENTRY\">>\"$FSTAB_FILE\"" || exit 1 -fi