0xTriboulet-evading_the_machine/classify.py

#!/usr/bin/env python3
"""
classify.py
-----------
Loads the trained logistic regression model and scaler,
extracts features from a target PE binary (stub function),
and performs malware classification.
"""

import torch
import torch.nn as nn
import joblib
import numpy as np
import sys
import os
import math         
import pefile        
import string       

# --------------------------------------------------
# Logistic Regression Model (same as in train.py)
# --------------------------------------------------
class LogisticRegression(nn.Module):
    def __init__(self, input_dim):
        super(LogisticRegression, self).__init__()
        self.linear = nn.Linear(input_dim, 1)

    def forward(self, x):
        return self.linear(x)  # raw logits


# --------------------------------------------------
# Feature Extraction Implementation
# --------------------------------------------------
def shannon_entropy(data: bytes) -> float:
    """Compute Shannon entropy of a byte sequence."""
    if len(data) == 0:
        return 0.0
    freq = np.bincount(np.frombuffer(data, dtype=np.uint8), minlength=256)
    probs = freq / len(data)
    probs = probs[probs > 0]
    return -np.sum(probs * np.log2(probs))


def extract_features(binary_path):
    """
    Extract PE features:
      - Weighted section entropy
      - Strings density
      - Log10 file size
    Returns: np.array([entropy, strings_density, log_size])
    """
    if not os.path.exists(binary_path):
        raise FileNotFoundError(f"File not found: {binary_path}")

    # --------------------------------------------------
    # File size
    # --------------------------------------------------
    file_size = os.path.getsize(binary_path)
    file_size_kb = max(file_size / 1024.0, 1e-6)  # avoid div-by-zero
    log_size = math.log10(file_size + 1)

    # --------------------------------------------------
    # Section entropy (size-weighted)
    # --------------------------------------------------
    try:
        pe = pefile.PE(binary_path, fast_load=True)
        section_entropies = []
        section_sizes = []
        for section in pe.sections:
            data = section.get_data()
            entropy = shannon_entropy(data)
            section_entropies.append(entropy)
            section_sizes.append(len(data))
        if section_sizes:
            weighted_entropy = np.average(section_entropies, weights=section_sizes)
        else:
            weighted_entropy = 0.0
    except Exception:
        weighted_entropy = 0.0

    # --------------------------------------------------
    # Strings density
    # --------------------------------------------------
    min_len = 4
    with open(binary_path, "rb") as f:
        raw_bytes = f.read()

    printable = set(bytes(string.printable, "ascii"))
    count_strings = 0
    current = bytearray()
    for b in raw_bytes:
        if b in printable and b not in b"\r\n\t":
            current.append(b)
        else:
            if len(current) >= min_len:
                count_strings += 1
            current = bytearray()
    if len(current) >= min_len:
        count_strings += 1

    strings_density = count_strings / file_size_kb

    # --------------------------------------------------
    # Return feature vector
    # --------------------------------------------------
    features = np.array([weighted_entropy, strings_density, log_size], dtype=np.float32)
    return features

# --------------------------------------------------
# Classification Logic
# --------------------------------------------------
def classify(binary_path):
    # Load scaler
    scaler = joblib.load("scaler.pkl")

    # Extract features
    raw_features = extract_features(binary_path).reshape(1, -1)
    features = scaler.transform(raw_features)

    # Load model
    input_dim = features.shape[1]
    model = LogisticRegression(input_dim)
    model.load_state_dict(torch.load("logistic_pe_model.pth", map_location="cpu"))
    model.eval()

    # Convert to tensor
    X = torch.tensor(features.astype("float32"))

    # Run model
    with torch.no_grad():
        logits = model(X)
        prob = torch.sigmoid(logits).item()

    # Threshold at 0.5
    label = 1 if prob >= 0.5 else 0
    verdict = "MALWARE" if label == 1 else "BENIGN"

    print(f"File: {binary_path}")
    print(f"Probability of malware: {prob:.4f}")
    print(f"Classification: {verdict}")


# --------------------------------------------------
# Entry Point
# --------------------------------------------------
if __name__ == "__main__":
    if len(sys.argv) != 2:
        print(f"Usage: python {sys.argv[0]} <target_binary>")
        sys.exit(1)

    target_binary = sys.argv[1]
    classify(target_binary)