Add files via upload

Engine/Properties/convert.py

@@ -1,4 +1,4 @@
-import re, os, sys, math, time, zlib, datetime, hashlib
+import re, os, sys, math, time, zlib, mmap, datetime, hashlib
 import ctypes, ctypes.wintypes, pefile, multiprocessing
 import numpy as np
 
@@ -113,6 +113,9 @@ def verify_signature(file_path):
 
 class FeatureExtractor:
     _STRING_PATTERN = re.compile(b'[\x20-\x7E]{5,}')
+    _STRING_FAST_PATH_BYTES = 4194304
+    _STRING_BATCH_BYTES = 4194304
+    _HISTOGRAM_BATCH_WINDOWS = 1024
     
     _BYTE_HIST_KEYS = [f'ByteHist_{i:02X}' for i in range(256)]
     _BYTE_ENT_KEYS = [f'ByteEnt_{i:02X}' for i in range(256)]
@@ -125,6 +128,49 @@ class FeatureExtractor:
     _C_LOG_C_TABLE = np.zeros(2049, dtype=np.float64)
     _C_LOG_C_TABLE[1:] = np.arange(1, 2049) * np.log2(np.arange(1, 2049))
 
+    _OPTIONAL_HEADER_FIELDS = [
+        'Magic', 'MajorLinkerVersion', 'MinorLinkerVersion', 'SizeOfCode',
+        'SizeOfInitializedData', 'SizeOfUninitializedData', 'AddressOfEntryPoint',
+        'BaseOfCode', 'ImageBase', 'SectionAlignment', 'FileAlignment',
+        'MajorOperatingSystemVersion', 'MinorOperatingSystemVersion',
+        'MajorImageVersion', 'MinorImageVersion', 'MajorSubsystemVersion',
+        'MinorSubsystemVersion', 'SizeOfImage', 'SizeOfHeaders', 'CheckSum',
+        'Subsystem', 'DllCharacteristics', 'SizeOfStackReserve', 'SizeOfStackCommit',
+        'SizeOfHeapReserve', 'SizeOfHeapCommit', 'LoaderFlags', 'NumberOfRvaAndSizes'
+    ]
+
+    _SECTION_DEFAULTS = {
+        'SectionCount': 0.0,
+        'SectionMaxEntropy': 0.0,
+        'SectionMinEntropy': 0.0,
+        'SectionMeanEntropy': 0.0,
+        'SectionMaxRawSize': 0.0,
+        'SectionMinRawSize': 0.0,
+        'SectionMeanRawSize': 0.0,
+        'SectionMaxVSize': 0.0,
+        'SectionMinVSize': 0.0,
+        'SectionMeanVSize': 0.0,
+        'ExecutableSections': 0.0,
+        'WritableSections': 0.0,
+        'ReadableSections': 0.0,
+        'SectionException': 0.0
+    }
+
+    _CONDITIONAL_DEFAULTS = {
+        'IsDriver': 0.0,
+        'HasTlsCallbacks': 0.0,
+        'IsDebug': 0.0,
+        'IsPreRelease': 0.0,
+        'IsPatched': 0.0,
+        'IsPrivateBuild': 0.0,
+        'IsSpecialBuild': 0.0,
+        'ImportCount': 0.0,
+        'ImportFunctionCount': 0.0,
+        'ExportCount': 0.0,
+        'IconCount': 0.0,
+        'DebugCount': 0.0
+    }
+
     @staticmethod
     def _safe_float(val):
         try:
@@ -134,6 +180,18 @@ class FeatureExtractor:
         except Exception: 
             return 0.0
 
+    @classmethod
+    def _base_defaults(cls):
+        base = dict.fromkeys(cls._OPTIONAL_HEADER_FIELDS, 0.0)
+        base.update(cls._SECTION_DEFAULTS)
+        base.update(cls._CONDITIONAL_DEFAULTS)
+
+        for i in range(16):
+            base[f'DataDirectory_{i}_Size'] = 0.0
+            base[f'DataDirectory_{i}_VA'] = 0.0
+
+        return base
+
     @staticmethod
     def _calc_entropy(data):
         if not data:
@@ -151,25 +209,84 @@ class FeatureExtractor:
         res = {"StringCount": 0.0, "StringMeanLength": 0.0, "StringEntropy": 0.0}
         res.update(dict.fromkeys(cls._STRING_KEYS, 0.0))
 
-        matches = cls._STRING_PATTERN.findall(file_bytes)
-        if not matches:
-            return res
+        if len(file_bytes) <= cls._STRING_FAST_PATH_BYTES:
+            matches = cls._STRING_PATTERN.findall(file_bytes)
+            if not matches:
+                return res
 
-        combined_data = b''.join(matches)
-        count = len(matches)
-        total_len = len(combined_data)
+            combined_data = b''.join(matches)
+            count = len(matches)
+            total_len = len(combined_data)
+            global_counts = np.bincount(np.frombuffer(combined_data, dtype=np.uint8), minlength=256)
+        else:
+            count = 0
+            total_len = 0
+            global_counts = np.zeros(256, dtype=np.int64)
+            mv = memoryview(file_bytes)
+            pos = 0
+            size = len(file_bytes)
+
+            while pos < size:
+                chunk_end = min(pos + cls._STRING_BATCH_BYTES, size)
+                chunk_arr = np.frombuffer(mv[pos:chunk_end], dtype=np.uint8)
+                separators = (chunk_arr < 0x20) | (chunk_arr > 0x7E)
+
+                if separators.any():
+                    last_separator = len(separators) - 1 - int(np.argmax(separators[::-1]))
+                    segment_end = pos + last_separator + 1
+                    matches = cls._STRING_PATTERN.findall(mv[pos:segment_end])
+
+                    if matches:
+                        combined_data = b''.join(matches)
+                        count += len(matches)
+                        total_len += len(combined_data)
+                        global_counts += np.bincount(np.frombuffer(combined_data, dtype=np.uint8), minlength=256)
+
+                    pos = segment_end
+                    continue
+
+                pending_len = 0
+                pending_counts = np.zeros(256, dtype=np.int64)
+
+                while pos < size:
+                    chunk_end = min(pos + cls._STRING_BATCH_BYTES, size)
+                    chunk_arr = np.frombuffer(mv[pos:chunk_end], dtype=np.uint8)
+                    separators = (chunk_arr < 0x20) | (chunk_arr > 0x7E)
+
+                    if separators.any():
+                        first_separator = int(np.argmax(separators))
+                        if first_separator:
+                            pending_counts += np.bincount(chunk_arr[:first_separator], minlength=256)
+                            pending_len += first_separator
+                        pos += first_separator + 1
+                        break
+
+                    pending_counts += np.bincount(chunk_arr, minlength=256)
+                    pending_len += len(chunk_arr)
+                    pos = chunk_end
+
+                if pending_len >= 5:
+                    count += 1
+                    total_len += pending_len
+                    global_counts += pending_counts
+
+            if count == 0 or total_len == 0:
+                return res
 
         res["StringCount"] = float(count)
         res["StringMeanLength"] = float(total_len) / count
-        res["StringEntropy"] = cls._calc_entropy(combined_data)
 
-        arr = np.frombuffer(combined_data, dtype=np.uint8)
-        counts = np.bincount(arr, minlength=127)
-        char_counts = counts[0x20:0x7F]
+        counts_nonzero = global_counts[global_counts > 0]
+        if len(counts_nonzero) > 0:
+            res["StringEntropy"] = float(np.log2(total_len) - np.sum(counts_nonzero * np.log2(counts_nonzero)) / total_len)
+
+        valid_chars = global_counts[0x20:0x7F]
+        res.update(dict(zip(cls._STRING_KEYS, valid_chars.astype(np.float64) / total_len)))
 
-        res.update(dict(zip(cls._STRING_KEYS, char_counts.astype(np.float64) / total_len)))
         return res
 
+
+
     @classmethod
     def _extract_histograms(cls, file_bytes):
         arr = np.frombuffer(file_bytes, dtype=np.uint8)
@@ -177,56 +294,64 @@ class FeatureExtractor:
         res = {}
 
         if sz == 0:
+            res['FileEntropy'] = 0.0
             res.update(dict.fromkeys(cls._BYTE_HIST_KEYS, 0.0))
             res.update(dict.fromkeys(cls._BYTE_ENT_KEYS, 0.0))
             return res
 
-        byte_counts = np.bincount(arr, minlength=256)
-        byte_hist = byte_counts.astype(np.float64) / sz
-        res.update(dict(zip(cls._BYTE_HIST_KEYS, byte_hist)))
-
+        byte_counts = np.zeros(256, dtype=np.int64)
         ent_hist = np.zeros(256, dtype=np.float64)
         window = 2048
         step = 1024
 
         if sz < window:
-            w_counts = byte_counts
-            p = w_counts[w_counts > 0] / float(sz)
+            byte_counts += np.bincount(arr, minlength=256)
+            p = byte_counts[byte_counts > 0] / float(sz)
             entropy = float(-np.sum(p * np.log2(p)))
             ent_bin = min(int(entropy * 2.0), 15)
-            byte_bins = w_counts.reshape(16, 16).sum(axis=1)
+            byte_bins = byte_counts.reshape(16, 16).sum(axis=1)
             idx_start = ent_bin * 16
             ent_hist[idx_start : idx_start+16] += byte_bins
-
         else:
             num_windows = (sz - window) // step + 1
-            batch_windows = 10000 
-            
-            for w_start in range(0, num_windows, batch_windows):
-                w_end = min(w_start + batch_windows, num_windows)
+
+            for w_start in range(0, num_windows, cls._HISTOGRAM_BATCH_WINDOWS):
+                w_end = min(w_start + cls._HISTOGRAM_BATCH_WINDOWS, num_windows)
                 b_num = w_end - w_start
-                
+
                 start_byte = w_start * step
                 end_byte = start_byte + b_num * step + step
-                b_arr = arr[start_byte:end_byte]
-                
-                arr_2d = b_arr.reshape(b_num + 1, step)
-                offsets = (np.arange(b_num + 1, dtype=np.int32)[:, None] * 256)
-                flat_idx = (arr_2d.astype(np.int32) + offsets).ravel()
-                
-                chunk_counts = np.bincount(flat_idx, minlength=(b_num + 1) * 256).reshape(b_num + 1, 256)
+                encoded = arr[start_byte:end_byte].astype(np.int32).reshape(b_num + 1, step)
+                encoded += np.arange(b_num + 1, dtype=np.int32)[:, None] * 256
+
+                chunk_counts = np.bincount(encoded.ravel(), minlength=(b_num + 1) * 256).reshape(b_num + 1, 256)
+                del encoded
+
+                byte_counts += np.sum(chunk_counts[:-1], axis=0)
+                if w_end == num_windows:
+                    byte_counts += chunk_counts[-1]
+
                 window_counts = chunk_counts[:-1] + chunk_counts[1:]
-                
                 sum_c_log_c = np.sum(cls._C_LOG_C_TABLE[window_counts], axis=1)
                 entropies = 11.0 - sum_c_log_c / 2048.0
-                
+
                 ent_bins = (entropies * 2.0).astype(np.int32)
                 np.clip(ent_bins, 0, 15, out=ent_bins)
-                
+
                 byte_bins = window_counts.reshape(b_num, 16, 16).sum(axis=2)
                 flat_ent_bins = (ent_bins[:, None] * 16 + np.arange(16)).ravel()
                 np.add.at(ent_hist, flat_ent_bins, byte_bins.ravel())
 
+            covered_size = (num_windows + 1) * step
+            if covered_size < sz:
+                byte_counts += np.bincount(arr[covered_size:], minlength=256)
+
+        counts_nonzero = byte_counts[byte_counts > 0]
+        res['FileEntropy'] = float(np.log2(sz) - np.sum(counts_nonzero * np.log2(counts_nonzero)) / sz)
+
+        byte_hist = byte_counts.astype(np.float64) / sz
+        res.update(dict(zip(cls._BYTE_HIST_KEYS, byte_hist)))
+
         ent_sum = np.sum(ent_hist)
         if ent_sum > 0:
             ent_hist /= ent_sum
@@ -234,6 +359,7 @@ class FeatureExtractor:
         res.update(dict(zip(cls._BYTE_ENT_KEYS, ent_hist)))
         return res
 
+
     @classmethod
     def _extract_overlay_features(cls, pe, fsize):
         overlay_offset = pe.get_overlay_data_start_offset()
@@ -390,7 +516,7 @@ class FeatureExtractor:
         if fsize == 0:
             return None
 
-        base, dlls, apis, pe = {}, set(), set(), None
+        base, dlls, apis, pe = cls._base_defaults(), set(), set(), None
         try:
             pe = pefile.PE(data=file_bytes, fast_load=True)
             
@@ -408,7 +534,6 @@ class FeatureExtractor:
                 pass
             
             base['TrustSigned'] = float(verify_signature(file_path))
-            base['FileEntropy'] = cls._calc_entropy(file_bytes)
             base['FileSize'] = float(fsize)
             
             fh = pe.FILE_HEADER
@@ -613,20 +738,22 @@ def _extract_worker_wrapper(args):
     return _extract_worker(*args)
 
 def _extract_worker(file_path, label):
+    mm = None
+    f = None
+
     try:
         fsize = os.path.getsize(file_path)
         if fsize == 0 or fsize > MAX_FILE_SIZE:
             return 'error', None
 
-        with open(file_path, "rb") as f:
-            file_bytes = f.read()
-
-        sha256 = hashlib.sha256(file_bytes).hexdigest()
+        f = open(file_path, "rb")
+        mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
+        sha256 = hashlib.sha256(mm).hexdigest()
         
         if sha256 in GLOBAL_HASHES:
             return 'skipped', None
 
-        res = FeatureExtractor.extract(file_bytes, file_path, fsize)
+        res = FeatureExtractor.extract(mm, file_path, fsize)
         if res:
             final_res = {"Label": label, "FileHash": sha256, **res}
             return 'success', (sha256, JSON_DUMPS(final_res) + "\n")
@@ -634,6 +761,11 @@ def _extract_worker(file_path, label):
         return 'error', None
     except Exception:
         return 'error', None
+    finally:
+        if mm:
+            mm.close()
+        if f:
+            f.close()
 
 ####################################################################################################
 

Engine/Properties/scanner.py

@@ -1,4 +1,4 @@
-import os, re, sys, math, json, zlib, pefile, datetime
+import os, re, sys, math, json, zlib, mmap, pefile, datetime
 import ctypes, ctypes.wintypes
 import numpy as np
 import onnxruntime as ort
@@ -99,6 +99,9 @@ class WinTrust:
 
 class FeatureExtractor:
     _STRING_PATTERN = re.compile(b'[\x20-\x7E]{5,}')
+    _STRING_FAST_PATH_BYTES = 4194304
+    _STRING_BATCH_BYTES = 4194304
+    _HISTOGRAM_BATCH_WINDOWS = 1024
     
     _BYTE_HIST_KEYS = [f'ByteHist_{i:02X}' for i in range(256)]
     _BYTE_ENT_KEYS = [f'ByteEnt_{i:02X}' for i in range(256)]
@@ -111,6 +114,49 @@ class FeatureExtractor:
     _C_LOG_C_TABLE = np.zeros(2049, dtype=np.float64)
     _C_LOG_C_TABLE[1:] = np.arange(1, 2049) * np.log2(np.arange(1, 2049))
 
+    _OPTIONAL_HEADER_FIELDS = [
+        'Magic', 'MajorLinkerVersion', 'MinorLinkerVersion', 'SizeOfCode',
+        'SizeOfInitializedData', 'SizeOfUninitializedData', 'AddressOfEntryPoint',
+        'BaseOfCode', 'ImageBase', 'SectionAlignment', 'FileAlignment',
+        'MajorOperatingSystemVersion', 'MinorOperatingSystemVersion',
+        'MajorImageVersion', 'MinorImageVersion', 'MajorSubsystemVersion',
+        'MinorSubsystemVersion', 'SizeOfImage', 'SizeOfHeaders', 'CheckSum',
+        'Subsystem', 'DllCharacteristics', 'SizeOfStackReserve', 'SizeOfStackCommit',
+        'SizeOfHeapReserve', 'SizeOfHeapCommit', 'LoaderFlags', 'NumberOfRvaAndSizes'
+    ]
+
+    _SECTION_DEFAULTS = {
+        'SectionCount': 0.0,
+        'SectionMaxEntropy': 0.0,
+        'SectionMinEntropy': 0.0,
+        'SectionMeanEntropy': 0.0,
+        'SectionMaxRawSize': 0.0,
+        'SectionMinRawSize': 0.0,
+        'SectionMeanRawSize': 0.0,
+        'SectionMaxVSize': 0.0,
+        'SectionMinVSize': 0.0,
+        'SectionMeanVSize': 0.0,
+        'ExecutableSections': 0.0,
+        'WritableSections': 0.0,
+        'ReadableSections': 0.0,
+        'SectionException': 0.0
+    }
+
+    _CONDITIONAL_DEFAULTS = {
+        'IsDriver': 0.0,
+        'HasTlsCallbacks': 0.0,
+        'IsDebug': 0.0,
+        'IsPreRelease': 0.0,
+        'IsPatched': 0.0,
+        'IsPrivateBuild': 0.0,
+        'IsSpecialBuild': 0.0,
+        'ImportCount': 0.0,
+        'ImportFunctionCount': 0.0,
+        'ExportCount': 0.0,
+        'IconCount': 0.0,
+        'DebugCount': 0.0
+    }
+
     @staticmethod
     def _safe_float(val):
         try:
@@ -120,6 +166,18 @@ class FeatureExtractor:
         except Exception: 
             return 0.0
 
+    @classmethod
+    def _base_defaults(cls):
+        base = dict.fromkeys(cls._OPTIONAL_HEADER_FIELDS, 0.0)
+        base.update(cls._SECTION_DEFAULTS)
+        base.update(cls._CONDITIONAL_DEFAULTS)
+
+        for i in range(16):
+            base[f'DataDirectory_{i}_Size'] = 0.0
+            base[f'DataDirectory_{i}_VA'] = 0.0
+
+        return base
+
     @staticmethod
     def _calc_entropy(data):
         if not data:
@@ -137,25 +195,84 @@ class FeatureExtractor:
         res = {"StringCount": 0.0, "StringMeanLength": 0.0, "StringEntropy": 0.0}
         res.update(dict.fromkeys(cls._STRING_KEYS, 0.0))
 
-        matches = cls._STRING_PATTERN.findall(file_bytes)
-        if not matches:
-            return res
+        if len(file_bytes) <= cls._STRING_FAST_PATH_BYTES:
+            matches = cls._STRING_PATTERN.findall(file_bytes)
+            if not matches:
+                return res
 
-        combined_data = b''.join(matches)
-        count = len(matches)
-        total_len = len(combined_data)
+            combined_data = b''.join(matches)
+            count = len(matches)
+            total_len = len(combined_data)
+            global_counts = np.bincount(np.frombuffer(combined_data, dtype=np.uint8), minlength=256)
+        else:
+            count = 0
+            total_len = 0
+            global_counts = np.zeros(256, dtype=np.int64)
+            mv = memoryview(file_bytes)
+            pos = 0
+            size = len(file_bytes)
+
+            while pos < size:
+                chunk_end = min(pos + cls._STRING_BATCH_BYTES, size)
+                chunk_arr = np.frombuffer(mv[pos:chunk_end], dtype=np.uint8)
+                separators = (chunk_arr < 0x20) | (chunk_arr > 0x7E)
+
+                if separators.any():
+                    last_separator = len(separators) - 1 - int(np.argmax(separators[::-1]))
+                    segment_end = pos + last_separator + 1
+                    matches = cls._STRING_PATTERN.findall(mv[pos:segment_end])
+
+                    if matches:
+                        combined_data = b''.join(matches)
+                        count += len(matches)
+                        total_len += len(combined_data)
+                        global_counts += np.bincount(np.frombuffer(combined_data, dtype=np.uint8), minlength=256)
+
+                    pos = segment_end
+                    continue
+
+                pending_len = 0
+                pending_counts = np.zeros(256, dtype=np.int64)
+
+                while pos < size:
+                    chunk_end = min(pos + cls._STRING_BATCH_BYTES, size)
+                    chunk_arr = np.frombuffer(mv[pos:chunk_end], dtype=np.uint8)
+                    separators = (chunk_arr < 0x20) | (chunk_arr > 0x7E)
+
+                    if separators.any():
+                        first_separator = int(np.argmax(separators))
+                        if first_separator:
+                            pending_counts += np.bincount(chunk_arr[:first_separator], minlength=256)
+                            pending_len += first_separator
+                        pos += first_separator + 1
+                        break
+
+                    pending_counts += np.bincount(chunk_arr, minlength=256)
+                    pending_len += len(chunk_arr)
+                    pos = chunk_end
+
+                if pending_len >= 5:
+                    count += 1
+                    total_len += pending_len
+                    global_counts += pending_counts
+
+            if count == 0 or total_len == 0:
+                return res
 
         res["StringCount"] = float(count)
         res["StringMeanLength"] = float(total_len) / count
-        res["StringEntropy"] = cls._calc_entropy(combined_data)
 
-        arr = np.frombuffer(combined_data, dtype=np.uint8)
-        counts = np.bincount(arr, minlength=127)
-        char_counts = counts[0x20:0x7F]
+        counts_nonzero = global_counts[global_counts > 0]
+        if len(counts_nonzero) > 0:
+            res["StringEntropy"] = float(np.log2(total_len) - np.sum(counts_nonzero * np.log2(counts_nonzero)) / total_len)
+
+        valid_chars = global_counts[0x20:0x7F]
+        res.update(dict(zip(cls._STRING_KEYS, valid_chars.astype(np.float64) / total_len)))
 
-        res.update(dict(zip(cls._STRING_KEYS, char_counts.astype(np.float64) / total_len)))
         return res
 
+
+
     @classmethod
     def _extract_histograms(cls, file_bytes):
         arr = np.frombuffer(file_bytes, dtype=np.uint8)
@@ -163,55 +280,64 @@ class FeatureExtractor:
         res = {}
 
         if sz == 0:
+            res['FileEntropy'] = 0.0
             res.update(dict.fromkeys(cls._BYTE_HIST_KEYS, 0.0))
             res.update(dict.fromkeys(cls._BYTE_ENT_KEYS, 0.0))
             return res
 
-        byte_counts = np.bincount(arr, minlength=256)
-        byte_hist = byte_counts.astype(np.float64) / sz
-        res.update(dict(zip(cls._BYTE_HIST_KEYS, byte_hist)))
-
+        byte_counts = np.zeros(256, dtype=np.int64)
         ent_hist = np.zeros(256, dtype=np.float64)
         window = 2048
         step = 1024
 
         if sz < window:
-            w_counts = byte_counts
-            p = w_counts[w_counts > 0] / float(sz)
+            byte_counts += np.bincount(arr, minlength=256)
+            p = byte_counts[byte_counts > 0] / float(sz)
             entropy = float(-np.sum(p * np.log2(p)))
             ent_bin = min(int(entropy * 2.0), 15)
-            byte_bins = w_counts.reshape(16, 16).sum(axis=1)
+            byte_bins = byte_counts.reshape(16, 16).sum(axis=1)
             idx_start = ent_bin * 16
             ent_hist[idx_start : idx_start+16] += byte_bins
         else:
             num_windows = (sz - window) // step + 1
-            batch_windows = 10000 
-            
-            for w_start in range(0, num_windows, batch_windows):
-                w_end = min(w_start + batch_windows, num_windows)
+
+            for w_start in range(0, num_windows, cls._HISTOGRAM_BATCH_WINDOWS):
+                w_end = min(w_start + cls._HISTOGRAM_BATCH_WINDOWS, num_windows)
                 b_num = w_end - w_start
-                
+
                 start_byte = w_start * step
                 end_byte = start_byte + b_num * step + step
-                b_arr = arr[start_byte:end_byte]
-                
-                arr_2d = b_arr.reshape(b_num + 1, step)
-                offsets = (np.arange(b_num + 1, dtype=np.int32)[:, None] * 256)
-                flat_idx = (arr_2d.astype(np.int32) + offsets).ravel()
-                
-                chunk_counts = np.bincount(flat_idx, minlength=(b_num + 1) * 256).reshape(b_num + 1, 256)
+                encoded = arr[start_byte:end_byte].astype(np.int32).reshape(b_num + 1, step)
+                encoded += np.arange(b_num + 1, dtype=np.int32)[:, None] * 256
+
+                chunk_counts = np.bincount(encoded.ravel(), minlength=(b_num + 1) * 256).reshape(b_num + 1, 256)
+                del encoded
+
+                byte_counts += np.sum(chunk_counts[:-1], axis=0)
+                if w_end == num_windows:
+                    byte_counts += chunk_counts[-1]
+
                 window_counts = chunk_counts[:-1] + chunk_counts[1:]
-                
                 sum_c_log_c = np.sum(cls._C_LOG_C_TABLE[window_counts], axis=1)
                 entropies = 11.0 - sum_c_log_c / 2048.0
-                
+
                 ent_bins = (entropies * 2.0).astype(np.int32)
                 np.clip(ent_bins, 0, 15, out=ent_bins)
-                
+
                 byte_bins = window_counts.reshape(b_num, 16, 16).sum(axis=2)
                 flat_ent_bins = (ent_bins[:, None] * 16 + np.arange(16)).ravel()
                 np.add.at(ent_hist, flat_ent_bins, byte_bins.ravel())
 
+            covered_size = (num_windows + 1) * step
+            if covered_size < sz:
+                byte_counts += np.bincount(arr[covered_size:], minlength=256)
+
+        counts_nonzero = byte_counts[byte_counts > 0]
+        res['FileEntropy'] = float(np.log2(sz) - np.sum(counts_nonzero * np.log2(counts_nonzero)) / sz)
+
+        byte_hist = byte_counts.astype(np.float64) / sz
+        res.update(dict(zip(cls._BYTE_HIST_KEYS, byte_hist)))
+
         ent_sum = np.sum(ent_hist)
         if ent_sum > 0:
             ent_hist /= ent_sum
@@ -219,6 +345,7 @@ class FeatureExtractor:
         res.update(dict(zip(cls._BYTE_ENT_KEYS, ent_hist)))
         return res
 
+
     @classmethod
     def _extract_overlay_features(cls, pe, fsize):
         overlay_offset = pe.get_overlay_data_start_offset()
@@ -372,15 +499,19 @@ class FeatureExtractor:
     @classmethod
     def extract(cls, file_path):
         pe = None
+        mm = None
+        source_file = None
+
         try:
             fsize = os.path.getsize(file_path)
             if fsize == 0 or fsize > MAX_FILE_SIZE: 
                 return None
 
-            with open(file_path, "rb") as f:
-                file_bytes = f.read()
+            source_file = open(file_path, "rb")
+            mm = mmap.mmap(source_file.fileno(), 0, access=mmap.ACCESS_READ)
+            file_bytes = mm
 
-            base, dlls, apis = {}, set(), set()
+            base, dlls, apis = cls._base_defaults(), set(), set()
             pe = pefile.PE(data=file_bytes, fast_load=True)
 
             try:
@@ -397,7 +528,6 @@ class FeatureExtractor:
                 pass
             
             base['TrustSigned'] = float(WinTrust.verify(file_path))
-            base['FileEntropy'] = cls._calc_entropy(file_bytes)
             base['FileSize'] = float(fsize)
             
             fh = pe.FILE_HEADER
@@ -420,7 +550,8 @@ class FeatureExtractor:
             if hasattr(pe, 'OPTIONAL_HEADER'):
                 op = pe.OPTIONAL_HEADER
                 fields = ['Magic', 'MajorLinkerVersion', 'MinorLinkerVersion', 'SizeOfCode', 'SizeOfInitializedData', 'SizeOfUninitializedData', 'AddressOfEntryPoint', 'BaseOfCode', 'ImageBase', 'SectionAlignment', 'FileAlignment', 'MajorOperatingSystemVersion', 'MinorOperatingSystemVersion', 'MajorImageVersion', 'MinorImageVersion', 'MajorSubsystemVersion', 'MinorSubsystemVersion', 'SizeOfImage', 'SizeOfHeaders', 'CheckSum', 'Subsystem', 'DllCharacteristics', 'SizeOfStackReserve', 'SizeOfStackCommit', 'SizeOfHeapReserve', 'SizeOfHeapCommit', 'LoaderFlags', 'NumberOfRvaAndSizes']
-                for f in fields: base[f] = getattr(op, f, 0)
+                for field_name in fields:
+                    base[field_name] = getattr(op, field_name, 0)
                 base['IsDriver'] = 1.0 if base.get('Subsystem') == 1 else 0.0
 
                 if hasattr(op, 'DATA_DIRECTORY'):
@@ -571,11 +702,24 @@ class FeatureExtractor:
 
             return {"Base": base, "DLLs": list(dlls), "APIs": list(apis)}
 
-        except Exception:
+        except pefile.PEFormatError:
             return None
         finally:
-            if pe:
-                pe.close()
+            if pe is not None:
+                try:
+                    pe.close()
+                except Exception:
+                    pass
+            if mm is not None:
+                try:
+                    mm.close()
+                except Exception:
+                    pass
+            if source_file is not None:
+                try:
+                    source_file.close()
+                except Exception:
+                    pass
 
 ####################################################################################################
 

Engine/Properties/train_server.py

@@ -58,15 +58,14 @@ def analyze_schema(jsonl_path):
         for line in f:
             num_lines += 1
 
-            if not base_keys:
-                try:
-                    data = JSON_LOADS(line)
-                    base_data = data.get('Base')
-                    if base_data:
-                        base_keys.update(base_data.keys())
+            try:
+                data = JSON_LOADS(line)
+                base_data = data.get('Base')
+                if base_data:
+                    base_keys.update(base_data.keys())
 
-                except Exception:
-                    pass
+            except Exception:
+                pass
 
     base_schema = sorted([k for k in base_keys if k.lower().strip() not in EXCLUDE_FEATURES])
     dll_schema = [f"DllHash_{i:03d}" for i in range(DLL_HASH_DIM)]

Engine/Properties/train_user.py

@@ -58,15 +58,14 @@ def analyze_schema(jsonl_path):
         for line in f:
             num_lines += 1
 
-            if not base_keys:
-                try:
-                    data = JSON_LOADS(line)
-                    base_data = data.get('Base')
-                    if base_data:
-                        base_keys.update(base_data.keys())
+            try:
+                data = JSON_LOADS(line)
+                base_data = data.get('Base')
+                if base_data:
+                    base_keys.update(base_data.keys())
 
-                except Exception:
-                    pass
+            except Exception:
+                pass
 
     base_schema = sorted([k for k in base_keys if k.lower().strip() not in EXCLUDE_FEATURES])
     dll_schema = [f"DllHash_{i:03d}" for i in range(DLL_HASH_DIM)]