src/predict.py

import lightgbm as lgb
import numpy as np
import ember
import sys
import json

MODEL_PATH = "data/ember_model_2018.txt"

def predict_pe(pe_path):
    """Use EMBER's feature extractor + pre-trained model to classify a PE file."""
    
    # Load model
    model = lgb.Booster(model_file=MODEL_PATH)

    # Extract features using EMBER's own extractor (2351-dim vector)
    with open(pe_path, "rb") as f:
        bytez = f.read()

    extractor = ember.PEFeatureExtractor(2)
    features = np.array(extractor.feature_vector(bytez), dtype=np.float32)

    # Score
    score = model.predict([features])[0]

    return {
        "file": pe_path,
        "malware_probability": round(float(score), 4),
        "verdict": "MALWARE" if score > 0.5 else "BENIGN",
        "confidence": f"{round(score * 100 if score > 0.5 else (1 - score) * 100, 1)}%"
    }


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python predict.py <path_to_exe>")
        sys.exit(1)

    result = predict_pe(sys.argv[1])
    print(json.dumps(result, indent=2))