From 586289bc763afab3c188c94cf48310c6c6788f00 Mon Sep 17 00:00:00 2001 From: Dobin Date: Sun, 21 Apr 2024 15:45:39 +0100 Subject: [PATCH] refactor: syntax improvement for dllresolver --- pe/dllresolver.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/pe/dllresolver.py b/pe/dllresolver.py index 5790b6c..2909d62 100644 --- a/pe/dllresolver.py +++ b/pe/dllresolver.py @@ -1,6 +1,7 @@ import sys import ctypes import os +from typing import List from pe.superpe import SuperPe @@ -12,13 +13,14 @@ class DllResolve(): self.path_res = path_res -def all_dll_exist(superpe): +def all_dll_exist(superpe) -> bool: for dll_name in superpe.get_iat_entries(): if not check_dll_availability(dll_name): return False return True -def unresolved_dlls(superpe): + +def unresolved_dlls(superpe) -> List[str]: res = [] for dll_name in superpe.get_iat_entries(): if not check_dll_availability(dll_name): @@ -26,20 +28,21 @@ def unresolved_dlls(superpe): return res -def resolve_dlls(superpe): +def resolve_dlls(superpe) -> List[DllResolve]: res = [] for dll_name in superpe.get_iat_entries(): res.append(resolve_dll(dll_name)) return res -def resolve_dll(dllname): +def resolve_dll(dllname) -> DllResolve: cdll_res = check_dll_availability(dllname) path_res = search_for_dll(dllname) return DllResolve(dllname, cdll_res, path_res) -def check_dll_availability(dll_name): +def check_dll_availability(dll_name) -> bool: + """Check if a DLL is available for loading by attempting to load it with ctypes.""" try: _ = ctypes.CDLL(dll_name) return True @@ -47,7 +50,8 @@ def check_dll_availability(dll_name): return False -def search_for_dll(dll_name): +def search_for_dll(dll_name) -> str: + """Search for a DLL in the system directories and PATH.""" paths = [ os.getcwd(), # Current directory os.environ.get('SYSTEMROOT', '') + '\\System32', # System directory