diff --git a/main.py b/main.py index 8b13789..d2d9823 100644 --- a/main.py +++ b/main.py @@ -1 +1,103 @@ +#!/usr/bin/env python3 + +import subprocess +import os +import re +import argparse + +SWAPFILE = "/swapfile" +FSTAB = "/etc/fstab" +GRUB_CONF = "/etc/default/grub" +MKINITCPIO = "/etc/mkinitcpio.conf" + +def run(cmd, capture=False): + if capture: + return subprocess.check_output(cmd, shell=True, text=True).strip() + subprocess.run(cmd, shell=True, check=True) + +def create_swapfile(size_gb): + print(f"[+] Creating {size_gb}G swapfile...") + run(f"fallocate -l {size_gb}G {SWAPFILE}") + run(f"chmod 600 {SWAPFILE}") + run(f"mkswap {SWAPFILE}") + run(f"swapon {SWAPFILE}") + +def update_fstab(): + print("[+] Ensuring swapfile is in /etc/fstab...") + with open(FSTAB, "r") as f: + if SWAPFILE in f.read(): + print("[-] Swapfile already in fstab.") + return + with open(FSTAB, "a") as f: + f.write(f"{SWAPFILE} none swap defaults 0 0\n") + +def get_swap_uuid(): + print("[+] Getting swap UUID...") + uuid = run(f"findmnt -no UUID -T {SWAPFILE}", capture=True) + return uuid + +def get_resume_offset(): + print("[+] Calculating resume_offset...") + out = run(f"filefrag -v {SWAPFILE}", capture=True) + match = re.search(r"^\s*0:\s+(\d+)", out, re.MULTILINE) + if match: + return match.group(1) + raise RuntimeError("Couldn't find resume offset.") + +def update_grub(uuid, offset): + print("[+] Updating GRUB_CMDLINE_LINUX_DEFAULT...") + with open(GRUB_CONF, "r") as f: + lines = f.readlines() + for i, line in enumerate(lines): + if line.startswith("GRUB_CMDLINE_LINUX_DEFAULT"): + line = re.sub(r'resume=UUID=\S+', '', line) + line = re.sub(r'resume_offset=\S+', '', line) + new = f'resume=UUID={uuid} resume_offset={offset}' + if '"' in line: + lines[i] = re.sub(r'"$', f' {new}"', line) + else: + lines[i] = line.strip() + f' {new}\n' + break + with open(GRUB_CONF, "w") as f: + f.writelines(lines) + run("update-grub") + +def update_mkinitcpio(): + print("[+] Ensuring resume hook in mkinitcpio.conf...") + with open(MKINITCPIO, "r") as f: + lines = f.readlines() + for i, line in enumerate(lines): + if line.startswith("HOOKS="): + if "resume" not in line: + line = line.strip().rstrip(")") + " resume)" + lines[i] = line + "\n" + break + with open(MKINITCPIO, "w") as f: + f.writelines(lines) + run("mkinitcpio -P") + +def main(): + if os.geteuid() != 0: + print("This script must be run as root.") + return + + parser = argparse.ArgumentParser(description="Configure hibernation with optional swapfile setup.") + parser.add_argument("--create-swapfile", action="store_true", help="Create and configure a swapfile") + parser.add_argument("--swap-size", type=int, default=32, help="Swapfile size in GB (default: 32)") + args = parser.parse_args() + + if args.create_swapfile: + create_swapfile(args.swap_size) + update_fstab() + + uuid = get_swap_uuid() + offset = get_resume_offset() + update_grub(uuid, offset) + update_mkinitcpio() + + print("\n✅ Hibernate setup complete. Please reboot your system:") + print(" sudo reboot") + +if __name__ == "__main__": + main()