scripts: extend IDA importer

scripts/ida_import.py

@@ -1,12 +1,25 @@
+import ctypes
 from pathlib import Path
-from typing import Optional
+from typing import Any, Dict, List, Optional
 
+import ida_typeinf
+import idaapi
 import idc
 
-PROGRESS_FILE = Path(__file__).parent / "docs" / "progress.txt"
+REPO_DIR = Path(__file__).parent.parent
+DOCS_DIR = REPO_DIR / "docs"
+TYPES_FILE = DOCS_DIR / "ida_types.h"
+VARIABLES_FILE = DOCS_DIR / "ida_variables.txt"
+PROGRESS_FILE = DOCS_DIR / "progress.txt"
 
 
 def to_int(source: str) -> Optional[int]:
+    source = source.strip()
+    if source.startswith("/*"):
+        source = source[2:]
+    if source.endswith("*/"):
+        source = source[:-2]
+    source = source.strip()
     if not source.replace("-", ""):
         return None
     if source.startswith(("0x", "0X")):
@@ -14,27 +27,66 @@ def to_int(source: str) -> Optional[int]:
     return int(source, 16)
 
 
-with PROGRESS_FILE.open("r", encoding="utf-8") as handle:
-    for line in handle:
-        line = line.strip()
-        if not line or line.startswith("#"):
-            continue
-        name, offset, size, flags = line.split(maxsplit=3)
-        offset = to_int(offset)
-        size = to_int(size)
+def import_types() -> None:
+    print(f"Importing type information from {TYPES_FILE}:")
+    error_count = idaapi.idc_parse_types(str(TYPES_FILE), idc.PT_FILE)
+    print(f"    done ({error_count} errors)")
 
-        # ignore inline functions
-        if offset is None:
-            continue
 
-        if not name.startswith("sub_"):
-            default_function_name = f"sub_{offset:x}"
-            print(f"renaming 0x{offset:08x} to {name}")
-            idc.set_name(offset, name)
+def import_variables() -> None:
+    print(f"Importing variables information from {VARIABLES_FILE}:")
+    with VARIABLES_FILE.open("r") as handle:
+        for line in handle:
+            line = line.strip()
+            if not line or line.startswith("#"):
+                continue
+            offset, decl = line.split(maxsplit=1)
+            offset = to_int(offset)
+            if True:
+                print(f"    renaming 0x{offset:08x} to {decl}")
 
-        if flags == "+":
-            idc.set_color(offset, idc.CIC_FUNC, 0xF0FFEE)
-        elif flags == "x":
-            idc.set_color(offset, idc.CIC_FUNC, 0xD8D8D8)
-        else:
-            idc.set_color(offset, idc.CIC_FUNC, idc.DEFCOLOR)
+                til = idaapi.get_idati()
+                ti = idaapi.tinfo_t()
+
+                name = idaapi.parse_decl(ti, til, decl, idaapi.PT_VAR)
+                if name.startswith("_"):
+                    name = name[1:]
+                if not name.startswith('dword_'):
+                    idaapi.set_name(offset, name)
+                idaapi.apply_tinfo(offset, ti, 0)
+    print("    done")
+
+
+def import_functions() -> None:
+    print(f"Importing types information from {PROGRESS_FILE}:")
+    with PROGRESS_FILE.open("r", encoding="utf-8") as handle:
+        for line in handle:
+            line = line.strip()
+            if not line or line.startswith("#"):
+                continue
+            name, offset, size, flags = line.split(maxsplit=3)
+            offset = to_int(offset)
+            size = to_int(size)
+
+            # ignore inline functions
+            if offset is None:
+                continue
+
+            old_name = idc.get_name(offset)
+            if old_name != name and not name.startswith("sub_"):
+                print(f"    renaming {old_name} to {name}")
+                idc.set_name(offset, name)
+
+            if flags == "+":
+                idc.set_color(offset, idc.CIC_FUNC, 0xF0FFEE)
+            elif flags == "x":
+                idc.set_color(offset, idc.CIC_FUNC, 0xD8D8D8)
+            else:
+                idc.set_color(offset, idc.CIC_FUNC, idc.DEFCOLOR)
+    print("    done")
+
+
+if __name__ == "__main__":
+    import_types()
+    import_variables()
+    import_functions()