Initial commit

2023-10-29 16:42:16 +01:00
commit 4ede0f9eb0
207 changed files with 34545 additions and 0 deletions
@@ -0,0 +1,89 @@
+import subprocess
+from enum import Enum
+
+class ASPM(Enum):
+    ASPM_DISABLED =   0x40
+    ASPM_L0s_ONLY =   0x41
+    ASPM_L1_ONLY =    0x42
+    ASPM_L1_AND_L0s = 0x43
+
+root_complex = "00:1c.4"
+endpoint = "05:00.0"
+value_to_set = ASPM.ASPM_L1_AND_L0s
+
+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:
+        print(f"Expected 256 bytes, only got {len(all_bytes)} bytes!")
+        print(f"Are you running this as root?")
+        exit()
+    return all_bytes
+
+def find_byte_to_patch(bytes, pos):
+    print(f"{hex(pos)} points to {hex(bytes[pos])}")
+    pos = bytes[pos]
+    print(f"Value at {hex(pos)} is {hex(bytes[pos])}")
+    if bytes[pos] != 0x10:
+        print("Value is not 0x10!")
+        print("Reading the next byte...")
+        pos += 0x1
+        return find_byte_to_patch(bytes, pos)
+    else:
+        print(f"Found the byte at: {hex(pos)}")
+        print("Adding 0x10 to the register...")
+        pos += 0x10
+        print(f"Final register reads: {hex(bytes[pos])}")
+        return pos
+
+def patch_byte(device, position, value):
+    subprocess.Popen([
+        "setpci",
+        "-s",
+        device,
+        f"{hex(position)}.B={hex(value)}"
+    ]).communicate()
+
+def patch_device(addr):
+    print(get_device_name(addr))
+    endpoint_bytes = read_all_bytes(addr)
+    byte_position_to_patch = find_byte_to_patch(endpoint_bytes, 0x34)
+
+    print(f"Byte to patch: {hex(byte_position_to_patch)}")
+    print(f"Byte is set to {hex(endpoint_bytes[byte_position_to_patch])}")
+    print(f"-> {ASPM(int(endpoint_bytes[byte_position_to_patch])).name}")
+    if int(endpoint_bytes[byte_position_to_patch]) != value_to_set.value:
+        print("Value doesn't match the one we want, setting it!")
+        patch_byte(addr, byte_position_to_patch, value_to_set.value)
+        new_bytes = read_all_bytes(addr)
+        print(f"Byte is set to {hex(new_bytes[byte_position_to_patch])}")
+        print(f"-> {ASPM(int(new_bytes[byte_position_to_patch])).name}")
+    else:
+        print("Nothing to patch!")
+
+def main():
+    patch_device(root_complex)
+    patch_device(endpoint)
+
+
+if __name__ == "__main__":
+    main()