logregtrain.py

#!/usr/bin/env python3

"""
train_logistic.py

Trains a single Logistic Regression model on the preprocessed data.
This is a simplified version of trainall.py.
Includes a toggleable --scale flag.

### MODIFIED to include:
- Feature importance analysis (model.coef_)
- Misclassified sample analysis (FPs and FNs)
"""

import os
import argparse
import joblib
import json
import numpy as np
import pandas as pd
from sklearn.metrics import classification_report, confusion_matrix, roc_auc_score

# --- Import Model ---
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler  # <-- IMPORT SCALER

# Directory where preprocessBest.py saved its output
PREPROCESSED_DIR = "data_separated_features32bit"

def evaluate_model(model, model_name, X_train, y_train, X_challenge, y_challenge):
    """
    Prints a full evaluation report for a trained model.
    (This function is unchanged from trainall.py)
    
    ### NEW: This function now returns the challenge set predictions ###
    """
    print(f"\n--- {model_name} Training Set Evaluation ---")
    y_pred_train = model.predict(X_train)
    print(classification_report(y_train, y_pred_train, target_names=['Benign', 'Malware']))
    
    y_pred_challenge = None # Initialize
    
    if X_challenge is not None and y_challenge is not None:
        if X_challenge.shape[0] == 0:
            print("\nChallenge data was loaded, but contains 0 samples. Skipping evaluation.")
            return None # ### NEW ###

        print(f"\n--- {model_name} Challenge Set Evaluation ---")
        y_pred_challenge = model.predict(X_challenge) # ### NEW: Capture predictions ###
        
        try:
            y_prob_challenge = model.predict_proba(X_challenge)[:, 1] # Get malware prob
        except Exception as e:
            print(f"Could not get predict_proba: {e}")
            y_prob_challenge = None

        print(classification_report(y_challenge, y_pred_challenge, target_names=['Benign', 'Malware']))

        try:
            tn, fp, fn, tp = confusion_matrix(y_challenge, y_pred_challenge, labels=[0, 1]).ravel()
        except ValueError:
            print("Could not generate confusion matrix for challenge set.")
            tn, fp, fn, tp = 0, 0, 0, 0

        fpr = fp / (fp + tn) if (fp + tn) > 0 else 0.0
        fnr = fn / (fn + tp) if (fn + tp) > 0 else 0.0
        
        if y_prob_challenge is not None:
            auc = roc_auc_score(y_challenge, y_prob_challenge) if len(np.unique(y_challenge)) > 1 else 0.0
        else:
            auc = 0.0

        print("\n--- CHALLENGE SET PERFORMANCE ---")
        print(f"   Target FNR: < 5.0%")
        print(f"   Target FPR: < 1.0%")
        print("---------------------------------")
        print(f"   Actual FNR: {fnr * 100:.4f}%")
        print(f"   Actual FPR: {fpr * 100:.4f}%")
        print(f"   AUC: {auc:.4f}")
        print("---------------------------------")
        
        if fnr < 0.05 and fpr < 0.01:
            print("SUCCESS: Model meets performance targets on the challenge set.")
        else:
            print("FAILURE: Model DOES NOT meet performance targets.")
    else:
        print("\nSkipping challenge set evaluation as no challenge samples were loaded.")

    return y_pred_challenge # ### NEW: Return predictions for analysis ###


def main():
    parser = argparse.ArgumentParser(description='Train a Logistic Regression model on preprocessed data')
    
    parser.add_argument('--preprocessed_dir', type=str, default=PREPROCESSED_DIR,
                        help="Directory containing X_train.npy, y_train.npy, etc.")
    parser.add_argument('--output_dir', type=str, default="trained_logreg",
                        help="Directory to save the trained model")
    
    parser.add_argument('--scale', action='store_true',
                        help="Apply StandardScaler to the data before training")
    
    args = parser.parse_args()
    
    os.makedirs(args.output_dir, exist_ok=True)

    # --- 1. Load Feature Names ---
    feature_names_path = os.path.join(args.preprocessed_dir, 'feature_names.txt')
    try:
        with open(feature_names_path, 'r') as f:
            # Load all 11 feature names
            all_feature_names = [line.split(': ')[1].strip() for line in f if ':' in line]
        print(f"Loaded {len(all_feature_names)} total features from {feature_names_path}")
    except FileNotFoundError:
        print(f"CRITICAL ERROR: 'feature_names.txt' not found in {args.preprocessed_dir}")
        return

    # --- 2. Load Data (as DataFrames) ---
    challenge_hashes = []
    try:
        X_train_np = np.load(os.path.join(args.preprocessed_dir, 'X_train.npy'))
        y_train = np.load(os.path.join(args.preprocessed_dir, 'y_train.npy'))
        X_challenge_np = np.load(os.path.join(args.preprocessed_dir, 'X_challenge.npy'))
        y_challenge = np.load(os.path.join(args.preprocessed_dir, 'y_challenge.npy'))
        
        challenge_hashes_path = os.path.join(args.preprocessed_dir, 'challenge_filenames.txt')
        with open(challenge_hashes_path, 'r') as f:
            challenge_hashes = [line.strip() for line in f.readlines()]
        print(f"Loaded {len(challenge_hashes)} challenge hashes.")
        
        # Load all 11 features into DataFrames
        X_train = pd.DataFrame(X_train_np, columns=all_feature_names)
        X_challenge = pd.DataFrame(X_challenge_np, columns=all_feature_names)
        
        print(f"Loaded {len(X_train)} training samples.")
        print(f"Loaded {len(X_challenge)} challenge samples.")
        
        if len(challenge_hashes) != len(X_challenge):
            print(f"CRITICAL WARNING: Mismatch between challenge samples ({len(X_challenge)}) and hashes ({len(challenge_hashes)}).")
            print("Misclassified sample analysis may be incorrect.")
            # Pad or truncate hashes to match X_challenge for stability, though it's not ideal
            if len(challenge_hashes) > len(X_challenge):
                challenge_hashes = challenge_hashes[:len(X_challenge)]
            else:
                challenge_hashes.extend(["MISSING_HASH"] * (len(X_challenge) - len(challenge_hashes)))
        
    except FileNotFoundError as e:
        print(f"CRITICAL ERROR: Data file not found in {args.preprocessed_dir}")
        print(f"File not found: {e.filename}")
        return
    except ValueError as e:
        print(f"CRITICAL ERROR: Mismatch between .npy file columns and feature_names.txt.")
        print(f"   Loaded {len(all_feature_names)} names, but .npy file has {X_train_np.shape[1]} columns.")
        print(f"   Error: {e}")
        print("   You may need to delete the 'data/preprocessed_challenge_optimized' directory and re-run preprocessBest.py")
        return


    # --- 3. DEFINE FEATURES TO DROP ---
    features_to_drop = [
        'mismatched_dll_characteristics',
        'import_richness',
    ]
    
    # Drop columns from DataFrames
    X_train = X_train.drop(columns=features_to_drop, errors='ignore')
    X_challenge = X_challenge.drop(columns=features_to_drop, errors='ignore')
    
    # Update our feature_names list to reflect the drop
    feature_names = [f for f in all_feature_names if f not in features_to_drop]
    
    print(f"Dropped {len(features_to_drop)} features. Using {len(feature_names)} features for training.")

    # --- 4. Define Model ---
    model_name = "logistic_regression"
    model = LogisticRegression(
        max_iter=10000, 
        C=0.0005,
        penalty='l2',
        solver='lbfgs',
    )

    # --- 5. Train and Evaluate Logistic Regression ---
    print(f"\n========================================================")
    print(f"         Training Model: {model_name.upper()}")
    print(f"========================================================")
    
    scaler = None # Initialize scaler
    
    try:
        if args.scale:
            print("Scaling data using StandardScaler...")
            scaler = StandardScaler()
            
            # Fit on X_train and transform X_train
            # We use the .values to get the numpy array, which is faster
            X_train_processed = scaler.fit_transform(X_train.values)
            
            # Transform X_challenge using the *same* scaler
            X_challenge_processed = scaler.transform(X_challenge.values)
            
            print("Scaling complete.")
        else:
            print("Scaling disabled. Training on raw feature values.")
            # Use the original data
            X_train_processed = X_train.values
            X_challenge_processed = X_challenge.values

        # We can fit directly on the processed (scaled or unscaled) data
        model.fit(X_train_processed, y_train)
                                              
        print(f"Training complete for {model_name}.")

        # --- Save Model ---
        model_path = os.path.join(args.output_dir, f"{model_name}_model.joblib")
        joblib.dump(model, model_path)
        print(f"Saved model to {model_path}")

        # --- NEW: Save Scaler if it was used ---
        if scaler:
            scaler_path = os.path.join(args.output_dir, f"{model_name}_scaler.joblib")
            joblib.dump(scaler, scaler_path)
            print(f"Saved scaler to {scaler_path}")
        # --- END Save Scaler ---
        
        # Save the new 9-feature order
        feature_path = os.path.join(args.output_dir, f"{model_name}_features.json")
        with open(feature_path, 'w') as f:
            json.dump(feature_names, f)
        print(f"Saved feature order to {feature_path}")

        # --- Evaluate ---
        y_pred_challenge = evaluate_model(model, model_name.upper(), X_train_processed, y_train, X_challenge_processed, y_challenge)

    except Exception as e:
        print(f"!!! ERROR training or evaluating {model_name}: {e}")
        
    print("\n--- Logistic Regression Trained and Evaluated ---")
    print(f"Model saved in: {args.output_dir}")
    
    print(f"\n========================================================")
    print(f"       Feature Importance: {model_name.upper()}")
    print(f"========================================================")
    
    try:
        # Get the coefficients from the trained model
        # For binary classification, model.coef_ is shape (1, n_features)
        coefficients = model.coef_[0]
        
        # Create a DataFrame to view them nicely
        coef_df = pd.DataFrame({
            'feature': feature_names,
            'coefficient': coefficients
        })
        
        # Add absolute value for sorting by importance
        coef_df['importance'] = coef_df['coefficient'].abs()
        coef_df = coef_df.sort_values(by='importance', ascending=False)
        
        print("Model Coefficients (from most to least important):")
        print(coef_df.to_string(index=False))
        
        print("\n--- Interpretation ---")
        print(" > High POSITIVE coefficient: Feature pushes model to predict MALWARE (1)")
        print(" > High NEGATIVE coefficient: Feature pushes model to predict BENIGN (0)")
        if not args.scale:
            print("\nWARNING: Scaling was DISABLED.")
            print("Coefficients are NOT directly comparable, as they are on different scales.")
            print("Re-run with --scale for a more accurate importance comparison.")
            
    except Exception as e:
        print(f"Could not analyze feature importance: {e}")



if __name__ == "__main__":
    main()