From af13b31ce3a7e98450571a9020c9ef34c5ab90d8 Mon Sep 17 00:00:00 2001
From: Kabbone <tobias@opel-online.de>
Date: Sun, 23 Mar 2025 11:13:08 +0100
Subject: [PATCH] add autoaspm and rotate screen button

---
 modules/hardware/autoaspm.py | 114 +++++++++++++++++++++++++++++++++++
 modules/wm/waybar.nix        |  14 ++++-
 2 files changed, 126 insertions(+), 2 deletions(-)
 create mode 100755 modules/hardware/autoaspm.py

diff --git a/modules/hardware/autoaspm.py b/modules/hardware/autoaspm.py
new file mode 100755
index 0000000..300aa7c
--- /dev/null
+++ b/modules/hardware/autoaspm.py
@@ -0,0 +1,114 @@
+#!/usr/bin/env python3
+
+# Original bash script by Luis R. Rodriguez
+# Re-written in Python by z8
+# Re-re-written to patch supported devices automatically by notthebee
+
+import re
+import subprocess
+import os
+import platform
+from enum import Enum
+
+class ASPM(Enum):
+    DISABLED = 0b00
+    L0s = 0b01
+    L1 = 0b10
+    L0sL1 = 0b11
+
+
+def run_prerequisites():
+    if platform.system() != "Linux":
+        raise OSError("This script only runs on Linux-based systems")
+    if not os.environ.get("SUDO_UID") and os.geteuid() != 0:
+        raise PermissionError("This script needs root privileges to run")
+    lspci_detected = subprocess.run(["which", "lspci"], stdout = subprocess.DEVNULL, stderr = subprocess.DEVNULL)
+    if lspci_detected.returncode > 0:
+        raise Exception("lspci not detected. Please install pciutils")
+    lspci_detected = subprocess.run(["which", "setpci"], stdout = subprocess.DEVNULL, stderr = subprocess.DEVNULL)
+    if lspci_detected.returncode > 0:
+        raise Exception("setpci not detected. Please install pciutils")
+
+
+def get_device_name(addr):
+    p = subprocess.Popen([
+        "lspci",
+        "-s",
+        addr,
+    ], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    return p.communicate()[0].splitlines()[0].decode()
+
+def read_all_bytes(device):
+    all_bytes = bytearray()
+    device_name = get_device_name(device)
+    p = subprocess.Popen([
+        "lspci",
+        "-s",
+        device,
+        "-xxx"
+    ], stdout= subprocess.PIPE, stderr=subprocess.PIPE)
+    ret = p.communicate()
+    ret = ret[0].decode()
+    for line in ret.splitlines():
+        if not device_name in line and ": " in line:
+            all_bytes.extend(bytearray.fromhex(line.split(": ")[1]))
+    if len(all_bytes) < 256:
+        exit()
+    return all_bytes
+
+def find_byte_to_patch(bytes, pos):
+    pos = bytes[pos]
+    if bytes[pos] != 0x10:
+        pos += 0x1
+        return find_byte_to_patch(bytes, pos)
+    else:
+        pos += 0x10
+        return pos
+
+def patch_byte(device, position, value):
+    subprocess.Popen([
+        "setpci",
+        "-s",
+        device,
+        f"{hex(position)}.B={hex(value)}"
+    ]).communicate()
+
+def patch_device(addr, aspm_value):
+    endpoint_bytes = read_all_bytes(addr)
+    byte_position_to_patch = find_byte_to_patch(endpoint_bytes, 0x34)
+    if int(endpoint_bytes[byte_position_to_patch]) & 0b11 != aspm_value.value:
+        patched_byte = int(endpoint_bytes[byte_position_to_patch])
+        patched_byte = patched_byte >> 2
+        patched_byte = patched_byte << 2
+        patched_byte = patched_byte | aspm_value.value
+
+        patch_byte(addr, byte_position_to_patch, patched_byte)
+        print(f"{addr}: Enabled ASPM {aspm_value.name}")
+    else:
+        print(f"{addr}: Already has ASPM {aspm_value.name} enabled")
+
+
+def list_supported_devices():
+    pcie_addr_regex = r"([0-9a-f]{2}:[0-9a-f]{2}.[0-9a-f])"
+    lspci = subprocess.run("lspci -vv", shell=True, capture_output=True).stdout
+    lspci_arr = re.split(pcie_addr_regex, str(lspci))[1:]
+    lspci_arr = [ x+y for x,y in zip(lspci_arr[0::2], lspci_arr[1::2]) ]
+
+    aspm_devices = {}
+    for dev in lspci_arr:
+        device_addr = re.findall(pcie_addr_regex, dev)[0]
+        if "ASPM" not in dev or "ASPM not supported" in dev:
+            continue
+        aspm_support = re.findall(r"ASPM (L[L0-1s ]*),", dev)
+        if aspm_support:
+            aspm_devices.update({device_addr: ASPM[aspm_support[0].replace(" ", "")]})
+    return aspm_devices
+
+
+def main():
+    run_prerequisites()
+    for device, aspm_mode in list_supported_devices().items():
+        patch_device(device, aspm_mode)
+
+if __name__ == "__main__":
+    main()
diff --git a/modules/wm/waybar.nix b/modules/wm/waybar.nix
index 1105e88..2f34395 100644
--- a/modules/wm/waybar.nix
+++ b/modules/wm/waybar.nix
@@ -54,7 +54,7 @@
             margin: 0 5px;
         }
         
-        #custom-vkeyboard, #custom-appkill {
+        #custom-vkeyboard, #custom-appkill, #custom-rotate {
             padding: 0px 50px;
             margin: 0px 0px;
             font-size: 20px;
@@ -135,13 +135,17 @@
             color: red;
             background-color: black;
         }
+        #custom-rotate {
+            color: white;
+            background-color: black;
+        }
       '';
       settings = [{
         layer = "bottom";
         position = "top";
         height = 22;
         tray = { spacing = 10; };
-        modules-center = [ "clock" "custom/appkill" ];
+        modules-center = [ "custom/rotate" "clock" "custom/appkill" ];
         modules-left = [ "sway/mode" "sway/workspaces" "sway/window" ];
         #modules-left = [ "wlr/workspaces" ];
         modules-right = [ "idle_inhibitor" "pulseaudio" "network" "cpu" "memory" "backlight" "temperature" "battery" "tray" ];
@@ -269,6 +273,12 @@
             on-click = "${pkgs.sway}/bin/swaymsg kill";
             tooltip = false;
         };
+        "custom/rotate" = {
+            format = "♽ ";
+            icon-size = 20;
+            on-click = "${pkgs.sway}/bin/swaymsg output eDP-1 transform 180 clockwise";
+            tooltip = false;
+        };
       }];
     };
   };