src/engine_a/inference.py

"""
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})"