From c5d195a2f76600c3df5e33a39cb5847449748185 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=81=8D=E5=85=AE=E6=83=9A=E5=85=AE?=
 <101191390+HIllya51@users.noreply.github.com>
Date: Tue, 25 Jun 2024 20:16:57 +0800
Subject: [PATCH] fix

---
 LunaTranslator/importanalysis.py | 10 ++++++++++
 LunaTranslator/retrieval.py      | 20 +++++++-------------
 2 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/LunaTranslator/importanalysis.py b/LunaTranslator/importanalysis.py
index 95c23a77..e1b92a5f 100644
--- a/LunaTranslator/importanalysis.py
+++ b/LunaTranslator/importanalysis.py
@@ -155,6 +155,7 @@ MEM_COMMIT = 0x00001000
 MEM_DECOMMIT = 0x00004000
 PAGE_READWRITE = 0x04
 IMAGE_NT_OPTIONAL_HDR32_MAGIC = 0x10B
+IMAGE_NT_OPTIONAL_HDR64_MAGIC = 0x20B
 IMAGE_DIRECTORY_ENTRY_IMPORT = 1
 
 
@@ -181,6 +182,8 @@ def Rva2Offset(rva, psh, pnt, IMAGE_NT_HEADERS):
             break
         pSeh += sizeof(IMAGE_SECTION_HEADER)
     pSeh = cast(pSeh, POINTER(IMAGE_SECTION_HEADER)).contents
+    if pSeh.VirtualAddress == 0 or pSeh.PointerToRawData == 0:
+        return -1
     return rva - pSeh.VirtualAddress + pSeh.PointerToRawData
 
 
@@ -200,6 +203,10 @@ def importanalysis(fname):
     if magic != IMAGE_NT_OPTIONAL_HDR32_MAGIC:
         ntheaders = cast(ntheaders_addr, POINTER(IMAGE_NT_HEADERS64)).contents
         IMAGE_NT_HEADERS = IMAGE_NT_HEADERS64
+        magic = ntheaders.OptionalHeader.Magic
+        if magic != IMAGE_NT_OPTIONAL_HDR64_MAGIC:
+            # 无效的文件
+            return []
     pSech = (
         ntheaders_addr
         + sizeof(DWORD)
@@ -223,6 +230,9 @@ def importanalysis(fname):
         offset = Rva2Offset(
             pImportDescriptor_data.Name, pSech, ntheaders_addr, IMAGE_NT_HEADERS
         )
+        if offset == -1:
+            # python3.dll，无导入
+            return []
         name = virtualpointer + offset
         collect.append((cast(name, c_char_p).value.decode(), offset))
         pImportDescriptor += sizeof(IMAGE_IMPORT_DESCRIPTOR)
diff --git a/LunaTranslator/retrieval.py b/LunaTranslator/retrieval.py
index e7b31c1c..adda03bb 100644
--- a/LunaTranslator/retrieval.py
+++ b/LunaTranslator/retrieval.py
@@ -2,6 +2,7 @@ import modulefinder, shutil, os, sys
 import builtins, platform
 import sys
 from importanalysis import importanalysis
+
 pyversion = platform.python_version()
 pyversion2 = "".join(pyversion.split(".")[:2])
 x86 = platform.architecture()[0] == "32bit"
@@ -30,16 +31,6 @@ print(py37Path)
 py37Pathwebview = os.path.join(py37Path, webviewappendix)
 
 
-def get_import_table(file_path):
-    pe = pefile.PE(file_path)
-    import_dlls = []
-    if hasattr(pe, "DIRECTORY_ENTRY_IMPORT"):
-        for entry in pe.DIRECTORY_ENTRY_IMPORT:
-            dll_name = entry.dll.decode("utf-8")
-            import_dlls.append(dll_name)
-    return import_dlls
-
-
 def get_dependencies(filename):
     saveopen = builtins.open
 
@@ -224,13 +215,14 @@ for f in collect:
     elif f.endswith(".exe") or f.endswith(".pyd") or f.endswith(".dll"):
         if f.endswith("Magpie.Core.exe"):
             continue
-        imports=importanalysis(f)
+        print(f)
+        imports = importanalysis(f)
         print(f, imports)
         if len(imports) == 0:
             continue
         with open(f, "rb") as ff:
             bs = bytearray(ff.read())
-        for _dll,offset in imports:
+        for _dll, offset in imports:
             if _dll.lower().startswith("api-ms-win-core"):
                 # 其实对于api-ms-win-core-winrt-XXX实际上是到ComBase.dll之类的，不过此项目中不包含这些
                 _target = "kernel32.dll"
@@ -241,7 +233,9 @@ for f in collect:
             _dll = _dll.encode()
             _target = _target.encode()
             # print(len(bs))
-            bs[offset : offset + len(_dll)] = _target + b"\0" * (len(_dll) - len(_target))
+            bs[offset : offset + len(_dll)] = _target + b"\0" * (
+                len(_dll) - len(_target)
+            )
             # print(len(bs))
         with open(f, "wb") as ff:
             ff.write(bs)