diff --git a/Engine/Properties/convert.py b/Engine/Properties/convert.py index 028b4c2..65f3174 100644 --- a/Engine/Properties/convert.py +++ b/Engine/Properties/convert.py @@ -233,6 +233,42 @@ class FeatureExtractor: cfg["HasSEHTable"] = 1.0 return cfg + @classmethod + def _extract_security_directory(cls, pe, file_bytes, fsize): + res = {"HasSignature": 0.0, "SignatureCount": 0.0} + if not hasattr(pe, 'OPTIONAL_HEADER') or not hasattr(pe.OPTIONAL_HEADER, 'DATA_DIRECTORY'): + return res + + sec_idx = 4 + directories = pe.OPTIONAL_HEADER.DATA_DIRECTORY + if len(directories) <= sec_idx: + return res + + sec_dir = directories[sec_idx] + if sec_dir.Size == 0 or sec_dir.VirtualAddress == 0: + return res + + res["HasSignature"] = 1.0 + offset = sec_dir.VirtualAddress + size = sec_dir.Size + + if offset + size > fsize or offset < 0 or size < 0: + return res + + count = 0 + curr = offset + end = offset + size + + while curr + 8 <= end: + dw_length = int.from_bytes(file_bytes[curr:curr+4], 'little') + if dw_length < 8: + break + count += 1 + curr += (dw_length + 7) & ~7 + + res["SignatureCount"] = float(count) + return res + @classmethod def extract(cls, file_bytes, file_path, fsize): if fsize == 0: return None @@ -241,6 +277,19 @@ class FeatureExtractor: try: pe = pefile.PE(data=file_bytes, fast_load=True) + try: + pe.parse_rich_header() + pe.parse_data_directories(directories=[ + pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_IMPORT'], + pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_EXPORT'], + pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_RESOURCE'], + pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_DEBUG'], + pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_TLS'], + pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG'] + ]) + except Exception: + pass + base['TrustSigned'] = verify_signature(file_path) base['FileEntropy'] = cls._calc_entropy(file_bytes) base['FileSize'] = fsize @@ -345,17 +394,6 @@ class FeatureExtractor: except Exception: continue - try: - pe.parse_data_directories(directories=[ - pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_IMPORT'], - pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_EXPORT'], - pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_RESOURCE'], - pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_DEBUG'], - pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_TLS'], - pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG'] - ]) - except Exception: pass - if hasattr(pe, 'DIRECTORY_ENTRY_TLS') and hasattr(pe.DIRECTORY_ENTRY_TLS, 'struct'): base['HasTlsCallbacks'] = 1.0 if getattr(pe.DIRECTORY_ENTRY_TLS.struct, 'AddressOfCallBacks', 0) != 0 else 0.0 @@ -399,6 +437,7 @@ class FeatureExtractor: base.update(cls._extract_ep_anomalies(pe)) base.update(cls._extract_advanced_resources(pe)) base.update(cls._extract_load_config(pe)) + base.update(cls._extract_security_directory(pe, file_bytes, fsize)) for k in base: base[k] = cls._safe_float(base[k]) return {"Base": base, "DLLs": list(dlls), "APIs": list(apis)} diff --git a/Engine/Properties/scanner.py b/Engine/Properties/scanner.py index 412b775..771e2d1 100644 --- a/Engine/Properties/scanner.py +++ b/Engine/Properties/scanner.py @@ -224,6 +224,42 @@ class FeatureExtractor: cfg["HasSEHTable"] = 1.0 return cfg + @classmethod + def _extract_security_directory(cls, pe, file_bytes, fsize): + res = {"HasSignature": 0.0, "SignatureCount": 0.0} + if not hasattr(pe, 'OPTIONAL_HEADER') or not hasattr(pe.OPTIONAL_HEADER, 'DATA_DIRECTORY'): + return res + + sec_idx = 4 + directories = pe.OPTIONAL_HEADER.DATA_DIRECTORY + if len(directories) <= sec_idx: + return res + + sec_dir = directories[sec_idx] + if sec_dir.Size == 0 or sec_dir.VirtualAddress == 0: + return res + + res["HasSignature"] = 1.0 + offset = sec_dir.VirtualAddress + size = sec_dir.Size + + if offset + size > fsize or offset < 0 or size < 0: + return res + + count = 0 + curr = offset + end = offset + size + + while curr + 8 <= end: + dw_length = int.from_bytes(file_bytes[curr:curr+4], 'little') + if dw_length < 8: + break + count += 1 + curr += (dw_length + 7) & ~7 + + res["SignatureCount"] = float(count) + return res + @classmethod def extract(cls, file_path): pe = None @@ -237,6 +273,19 @@ class FeatureExtractor: base, dlls, apis = {}, set(), set() pe = pefile.PE(data=file_bytes, fast_load=True) + + try: + pe.parse_rich_header() + pe.parse_data_directories(directories=[ + pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_IMPORT'], + pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_EXPORT'], + pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_RESOURCE'], + pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_DEBUG'], + pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_TLS'], + pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG'] + ]) + except Exception: + pass base['TrustSigned'] = float(WinTrust.verify(file_path)) base['FileEntropy'] = cls._calc_entropy(file_bytes) @@ -342,17 +391,6 @@ class FeatureExtractor: except Exception: continue - try: - pe.parse_data_directories(directories=[ - pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_IMPORT'], - pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_EXPORT'], - pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_RESOURCE'], - pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_DEBUG'], - pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_TLS'], - pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG'] - ]) - except Exception: pass - if hasattr(pe, 'DIRECTORY_ENTRY_TLS') and hasattr(pe.DIRECTORY_ENTRY_TLS, 'struct'): base['HasTlsCallbacks'] = 1.0 if getattr(pe.DIRECTORY_ENTRY_TLS.struct, 'AddressOfCallBacks', 0) != 0 else 0.0 @@ -396,6 +434,7 @@ class FeatureExtractor: base.update(cls._extract_ep_anomalies(pe)) base.update(cls._extract_advanced_resources(pe)) base.update(cls._extract_load_config(pe)) + base.update(cls._extract_security_directory(pe, file_bytes, fsize)) for k in base: base[k] = cls._safe_float(base[k]) return {"Base": base, "DLLs": list(dlls), "APIs": list(apis)}