""" Inference script for Engine A (EMBER). This module handles loading a pre-trained LightGBM structural classifier and parsing uploaded Windows executable binaries (.exe, .dll) using pefile to approximate EMBER features. It also supports loading mock JSON profiles. """ import os import pickle import pefile import numpy as np import warnings import json class EngineAInfer: """ Inference Engine for Static Structural Malware Detection. Attributes: model: The loaded LightGBM classifier. feature_names (dict): Mapping of extracted array indices to human-readable PE header names. """ def __init__(self, model_path="models/engine_a_model.pkl"): """ Initializes the inference engine by loading the specified LightGBM model. Args: model_path (str): Path to the pickled LightGBM model file. """ self.model = None if os.path.exists(model_path): with open(model_path, "rb") as f: self.model = pickle.load(f) else: print( f"Warning: {model_path} not found. Engine A will fail unless trained." ) # Human readable names for the custom extracted pefile features self.feature_names = { 0: "Machine Architecture", 1: "SizeOfOptionalHeader", 2: "Characteristics", 3: "MajorLinkerVersion", 4: "MinorLinkerVersion", 5: "SizeOfCode", 6: "SizeOfInitializedData", 7: "SizeOfUninitializedData", 8: "AddressOfEntryPoint", 9: "BaseOfCode", 10: "NumberOfSections", 11: "NumberOfImports", 12: "TotalImportedFunctions", 13: "NumberOfExports", } def extract_features(self, file_path): """ Approximates EMBER features using the pefile library. Parses basic PE headers, sections, imports, and exports, and pads the vector to the 2381 features expected by the EMBER dataset format. Args: file_path (str): Absolute or relative path to the executable file. Returns: np.ndarray: A 2D numpy array of shape (1, 2381) representing the features. """ features = np.zeros(2381, dtype=np.float32) try: pe = pefile.PE(file_path) # Basic header features features[0] = pe.FILE_HEADER.Machine features[1] = pe.FILE_HEADER.SizeOfOptionalHeader features[2] = pe.FILE_HEADER.Characteristics features[3] = pe.OPTIONAL_HEADER.MajorLinkerVersion features[4] = pe.OPTIONAL_HEADER.MinorLinkerVersion features[5] = pe.OPTIONAL_HEADER.SizeOfCode features[6] = pe.OPTIONAL_HEADER.SizeOfInitializedData features[7] = pe.OPTIONAL_HEADER.SizeOfUninitializedData features[8] = pe.OPTIONAL_HEADER.AddressOfEntryPoint features[9] = pe.OPTIONAL_HEADER.BaseOfCode # Number of sections features[10] = pe.FILE_HEADER.NumberOfSections # Imports if hasattr(pe, "DIRECTORY_ENTRY_IMPORT"): features[11] = len(pe.DIRECTORY_ENTRY_IMPORT) num_imports = sum( [len(entry.imports) for entry in pe.DIRECTORY_ENTRY_IMPORT] ) features[12] = num_imports # Exports if hasattr(pe, "DIRECTORY_ENTRY_EXPORT"): features[13] = len(pe.DIRECTORY_ENTRY_EXPORT.symbols) except Exception as e: print(f"PE parsing error (might not be a PE file): {e}") return features.reshape(1, -1) def predict(self, file_path): """ Predicts if a file is malicious based on its structural layout. Also calculates SHAP (SHapley Additive exPlanations) values to provide explainability regarding which structural features influenced the decision. Args: file_path (str): Path to the target file. Can be a Windows PE binary or a simulated mock malware profile (.json). Returns: dict: Contains 'is_malware' (bool), 'malware_prob' (float), and 'top_features' (list of tuples detailing top contributing factors). """ if self.model is None: raise Exception("Engine A model is not loaded. Train the model first.") # Handle mock JSON profile logic if file_path.endswith(".json"): try: with open(file_path, "r") as f: data = json.load(f) if data.get("is_mock_profile"): features = np.array(data["ember_features"]).reshape(1, -1) else: features = self.extract_features(file_path) except: features = self.extract_features(file_path) else: features = self.extract_features(file_path) # Suppress LGBM missing feature names warning with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) probs = self.model.predict_proba(features)[0] # Use LightGBM pred_contrib=True to get SHAP values for explainability contributions = self.model.predict(features, pred_contrib=True)[0] malware_prob = float(probs[1]) if len(probs) > 1 else float(probs[0]) # Explainability Logic: Exclude the last element (base expected value) feature_shap = contributions[:-1] # If the file is malicious, we want the most positive SHAP values. # If it's benign, we want the most negative SHAP values. is_malicious = malware_prob > 0.5 if is_malicious: # Sort descending for highest positive push top_indices = np.argsort(feature_shap)[::-1][:3] else: # Sort ascending for lowest negative push top_indices = np.argsort(feature_shap)[:3] top_features = [] for idx in top_indices: name = self._get_feature_name(idx) contrib = feature_shap[idx] top_features.append((name, float(contrib))) return { "is_malware": is_malicious, "malware_prob": malware_prob, "top_features": top_features, } def _get_feature_name(self, idx): """ Maps an EMBER numerical feature index to a human-readable category. EMBER extracts 2381 features using LIEF. Because the parquet file strips the column names, we rely on the known ranges of the EMBER feature specification to explain what the model is looking at. """ if idx in self.feature_names and idx <= 13: # Our custom exact mappings for the first 14 features return f"{self.feature_names[idx]} (Index {idx})" # Generalized EMBER feature mapping by vector location if 0 <= idx <= 255: return f"Raw Byte Histogram Analysis (Index {idx})" elif 256 <= idx <= 511: return f"Byte Entropy/Complexity Matrix (Index {idx})" elif 512 <= idx <= 615: return f"Embedded Strings Metadata [Paths/URLs/RegKeys] (Index {idx})" elif 616 <= idx <= 626: return f"General PE Structural Info (Index {idx})" elif 627 <= idx <= 688: return f"PE Header Anomaly (Index {idx})" elif 689 <= idx <= 944: return f"PE Sections Characteristics (Index {idx})" elif 945 <= idx <= 2224: return f"Imported Libraries/API Calls (Index {idx})" elif 2225 <= idx <= 2381: return f"Exported Functions/Data Directories (Index {idx})" else: return f"Deep Structural Feature (Index {idx})"