From 8473113ef9331d5869d9fcb261578c62d6b674d7 Mon Sep 17 00:00:00 2001
From: seth <seth@ske.moe>
Date: Mon, 17 May 2021 21:52:35 +0000
Subject: [PATCH] Add support for 64 bit PE in psmap

---
 bemani/utils/psmap.py | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/bemani/utils/psmap.py b/bemani/utils/psmap.py
index 78c802f..2b4c7af 100644
--- a/bemani/utils/psmap.py
+++ b/bemani/utils/psmap.py
@@ -43,12 +43,18 @@ def parse_psmap(data: bytes, offset: str, rootname: str) -> Node:
     saved_loc: List[int] = []
 
     while True:
-        chunk = data[base:(base + 16)]
-        base = base + 16
+        if hex(pe.FILE_HEADER.Machine) == '0x8664':  # 64 bit
+            chunk = data[base:(base + 24)]
+            base = base + 24
 
-        (nodetype, mandatory, outoffset, width, nameptr, defaultptr) = struct.unpack('<BBHIII', chunk)
+            (nodetype, mandatory, outoffset, width, nameptr, defaultptr) = struct.unpack('<BBHIQQ', chunk)
+        else:  # 32 bit
+            chunk = data[base:(base + 16)]
+            base = base + 16
 
-        if nodetype == 0xFF:
+            (nodetype, mandatory, outoffset, width, nameptr, defaultptr) = struct.unpack('<BBHIII', chunk)
+
+        if nodetype == 0xFF or nodetype == 0x00:  # if nodetype is 0 then we probably read garbage
             # End of nodes, see if we should exit
             if len(saved_root) == 0:
                 break
@@ -67,9 +73,7 @@ def parse_psmap(data: bytes, offset: str, rootname: str) -> Node:
         except ValueError:
             pass
 
-        if nodetype == 0x00:
-            raise Exception(f'Invalid node type 0x{nodetype:02x}')
-        elif nodetype == 0x01:
+        if nodetype == 0x01:
             # This is a void node, so we should handle by recursing
             node = Node.void(name)
             root.add_child(node)