Update main.py

2025-04-02 16:12:33 +02:00
parent 4e780944ca
commit 4de9da0068
1 changed files with 102 additions and 0 deletions
--- 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()