android_malware_preprocessing.py

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.feature_selection import SelectFromModel
import os
import pickle
import warnings
from datetime import datetime

# Ignore warnings
warnings.filterwarnings('ignore')

def preprocess_android_malware_data(input_file, output_dir='preprocessed_data', 
                                   test_size=0.2, random_state=42, 
                                   create_visualizations=True, perform_feature_selection=True):
    """
    Complete preprocessing pipeline for Android malware detection dataset.
    
    Args:
        input_file: Path to the cleaned CSV file
        output_dir: Directory to save preprocessed data and visualizations
        test_size: Proportion of the dataset to include in the test split
        random_state: Random seed for reproducibility
        create_visualizations: Whether to create and save visualizations
        perform_feature_selection: Whether to perform feature selection
        
    Returns:
        Dictionary with preprocessing results and statistics
    """
    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)
    
    # Create a subdirectory for visualizations
    viz_dir = os.path.join(output_dir, 'visualizations')
    os.makedirs(viz_dir, exist_ok=True)
    
    # Dictionary to store preprocessing results
    results = {
        'input_file': input_file,
        'output_dir': output_dir,
        'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        'statistics': {}
    }
    
    print(f"Starting preprocessing of {input_file}...")
    
    # Step 1: Load the data
    print("\nStep 1: Loading data...")
    data = pd.read_csv(input_file)
    results['statistics']['initial_shape'] = data.shape
    print(f"Dataset loaded with {data.shape[0]} rows and {data.shape[1]} columns")
    
    # Step 2: Check for and handle missing values
    print("\nStep 2: Handling missing values...")
    missing_values = data.isnull().sum()
    results['statistics']['missing_values'] = missing_values[missing_values > 0].to_dict()
    
    if missing_values.sum() > 0:
        print(f"Found {missing_values.sum()} missing values in {len(missing_values[missing_values > 0])} columns")
        # For numeric features, fill with median
        numeric_cols = data.select_dtypes(include=['float64', 'int64']).columns
        for col in numeric_cols:
            if data[col].isnull().sum() > 0:
                data[col] = data[col].fillna(data[col].median())
                print(f"  Filled missing values in {col} with median")
    else:
        print("No missing values found in the dataset")
    
    # Step 3: Check class distribution
    print("\nStep 3: Checking class distribution...")
    class_distribution = data['is_malware'].value_counts()
    results['statistics']['class_distribution'] = class_distribution.to_dict()
    
    print(f"Class distribution:")
    print(f"  Benign samples (0): {class_distribution.get(0, 0)}")
    print(f"  Malware samples (1): {class_distribution.get(1, 0)}")
    
    if create_visualizations:
        plt.figure(figsize=(8, 6))
        sns.countplot(x='is_malware', data=data)
        plt.title('Class Distribution')
        plt.xlabel('Class (0=Benign, 1=Malware)')
        plt.ylabel('Count')
        plt.tight_layout()
        plt.savefig(os.path.join(viz_dir, 'class_distribution.png'))
        plt.close()
    
    # Step 4: Handle outliers
    print("\nStep 4: Handling outliers...")
    
    # Identify numeric columns for outlier detection
    numeric_cols = data.select_dtypes(include=['float64', 'int64']).columns
    numeric_cols = [col for col in numeric_cols if col != 'is_malware']
    
    # Function to detect outliers using IQR method
    def detect_outliers(df, col):
        Q1 = df[col].quantile(0.25)
        Q3 = df[col].quantile(0.75)
        IQR = Q3 - Q1
        lower_bound = Q1 - 1.5 * IQR
        upper_bound = Q3 + 1.5 * IQR
        outliers = df[(df[col] < lower_bound) | (df[col] > upper_bound)][col]
        return outliers.shape[0], lower_bound, upper_bound
    
    outlier_stats = {}
    capped_columns = []
    
    for col in numeric_cols:
        outlier_count, lower_bound, upper_bound = detect_outliers(data, col)
        outlier_stats[col] = {
            'count': outlier_count,
            'percentage': outlier_count / data.shape[0] * 100,
            'lower_bound': lower_bound,
            'upper_bound': upper_bound
        }
        
        # Only cap if there are significant outliers
        if outlier_count > data.shape[0] * 0.05:  # If more than 5% are outliers
            original_min = data[col].min()
            original_max = data[col].max()
            data[col] = data[col].clip(lower=lower_bound, upper=upper_bound)
            capped_columns.append(col)
            print(f"  Capped outliers in {col}: {outlier_count} values ({outlier_count/data.shape[0]:.2%})")
            print(f"    Original range: [{original_min}, {original_max}]")
            print(f"    New range: [{data[col].min()}, {data[col].max()}]")
    
    results['statistics']['outlier_stats'] = outlier_stats
    results['statistics']['capped_columns'] = capped_columns
    
    # Step 5: Create derived features
    print("\nStep 5: Creating derived features...")
    
    # Dictionary to track which derived features were successfully created
    derived_features = {}
    
    # Try to create each derived feature, handling potential division by zero
    try:
        data['dangerous_perm_ratio'] = data['dangerous_permission_count'] / data['permission_count'].replace(0, 1)
        derived_features['dangerous_perm_ratio'] = True
    except Exception as e:
        derived_features['dangerous_perm_ratio'] = str(e)
    
    try:
        data['service_to_activity_ratio'] = data['service_count'] / data['activity_count'].replace(0, 1)
        derived_features['service_to_activity_ratio'] = True
    except Exception as e:
        derived_features['service_to_activity_ratio'] = str(e)
    
    try:
        data['receiver_to_activity_ratio'] = data['receiver_count'] / data['activity_count'].replace(0, 1)
        derived_features['receiver_to_activity_ratio'] = True
    except Exception as e:
        derived_features['receiver_to_activity_ratio'] = str(e)
    
    try:
        data['method_to_class_ratio'] = data['method_count'] / data['class_count'].replace(0, 1)
        derived_features['method_to_class_ratio'] = True
    except Exception as e:
        derived_features['method_to_class_ratio'] = str(e)
    
    try:
        data['api_to_method_ratio'] = data['api_call_count'] / data['method_count'].replace(0, 1)
        derived_features['api_to_method_ratio'] = True
    except Exception as e:
        derived_features['api_to_method_ratio'] = str(e)
    
    try:
        data['url_to_string_ratio'] = data['url_string_count'] / data['string_count'].replace(0, 1)
        derived_features['url_to_string_ratio'] = True
    except Exception as e:
        derived_features['url_to_string_ratio'] = str(e)
    
    try:
        data['base64_to_string_ratio'] = data['base64_string_count'] / data['string_count'].replace(0, 1)
        derived_features['base64_to_string_ratio'] = True
    except Exception as e:
        derived_features['base64_to_string_ratio'] = str(e)
    
    try:
        data['reflection_to_api_ratio'] = data['reflection_count'] / data['api_call_count'].replace(0, 1)
        derived_features['reflection_to_api_ratio'] = True
    except Exception as e:
        derived_features['reflection_to_api_ratio'] = str(e)
    
    try:
        data['suspicious_behavior_score'] = (
            data['obfuscation_score'] + 
            data.get('reflection_to_api_ratio', 0) + 
            data.get('dangerous_perm_ratio', 0)
        ) / 3
        derived_features['suspicious_behavior_score'] = True
    except Exception as e:
        derived_features['suspicious_behavior_score'] = str(e)
    
    # Create compound permission features
    try:
        data['sms_and_internet'] = ((data['perm_read_sms'] > 0) | (data['perm_send_sms'] > 0)) & (data['network_api_count'] > 0)
        data['sms_and_internet'] = data['sms_and_internet'].astype(int)
        derived_features['sms_and_internet'] = True
    except Exception as e:
        derived_features['sms_and_internet'] = str(e)
    
    try:
        data['phone_and_internet'] = (data['perm_read_phone_state'] > 0) & (data['network_api_count'] > 0)
        data['phone_and_internet'] = data['phone_and_internet'].astype(int)
        derived_features['phone_and_internet'] = True
    except Exception as e:
        derived_features['phone_and_internet'] = str(e)
    
    # Create code complexity metrics
    try:
        data['code_complexity'] = (
            data['method_count'] * 0.4 + 
            data['opcode_diversity'] * 0.3 + 
            data['string_count'] * 0.3
        ) / 1000  # Normalize
        derived_features['code_complexity'] = True
    except Exception as e:
        derived_features['code_complexity'] = str(e)
    
    # Log successfully created features
    successful_features = [f for f, status in derived_features.items() if status is True]
    print(f"  Created {len(successful_features)} new derived features:")
    for feature in successful_features:
        print(f"    - {feature}")
    
    results['statistics']['derived_features'] = derived_features
    results['statistics']['dataset_shape_after_feature_engineering'] = data.shape
    
    # Step 6: Analyze feature correlations
    print("\nStep 6: Analyzing feature correlations...")
    # Calculate correlation matrix
    corr_matrix = data.corr()
    
    if create_visualizations:
        # Plot correlation heatmap
        plt.figure(figsize=(20, 16))
        mask = np.triu(corr_matrix)
        sns.heatmap(corr_matrix, mask=mask, cmap='coolwarm', annot=False, vmin=-1, vmax=1, 
                   square=True, linewidths=0.5)
        plt.title('Feature Correlation Matrix')
        plt.tight_layout()
        plt.savefig(os.path.join(viz_dir, 'correlation_heatmap.png'))
        plt.close()
    
    # Identify highly correlated features
    high_corr_threshold = 0.8
    high_corr_pairs = []
    high_corr_features = set()
    
    for i in range(len(corr_matrix.columns)):
        for j in range(i):
            if abs(corr_matrix.iloc[i, j]) > high_corr_threshold:
                colname_i = corr_matrix.columns[i]
                colname_j = corr_matrix.columns[j]
                if colname_i != 'is_malware' and colname_j != 'is_malware':
                    high_corr_pairs.append((colname_i, colname_j, corr_matrix.iloc[i, j]))
                    high_corr_features.add(colname_i)
    
    if high_corr_pairs:
        print(f"  Found {len(high_corr_pairs)} highly correlated feature pairs (r > {high_corr_threshold}):")
        for col1, col2, corr in high_corr_pairs[:5]:  # Show first 5 pairs
            print(f"    - {col1} and {col2}: {corr:.2f}")
        if len(high_corr_pairs) > 5:
            print(f"    - ... and {len(high_corr_pairs) - 5} more")
    else:
        print("  No highly correlated feature pairs found")
    
    results['statistics']['high_correlation_pairs'] = [(col1, col2, corr) for col1, col2, corr in high_corr_pairs]
    results['statistics']['high_correlation_count'] = len(high_corr_pairs)
    
    # Step 7: Separate features and target
    print("\nStep 7: Separating features and target...")
    X = data.drop('is_malware', axis=1)
    y = data['is_malware']
    
    # Step 8: Standardize features
    print("\nStep 8: Standardizing features...")
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    print(f"  Features have been standardized")
    
    # Feature names for future reference
    feature_names = X.columns.tolist()
    
    # Step 9: Feature selection (if requested)
    if perform_feature_selection:
        print("\nStep 9: Performing feature selection...")
        # Use Random Forest for feature importance-based selection
        rf = RandomForestClassifier(n_estimators=100, random_state=random_state)
        rf.fit(X_scaled, y)
        
        # Get feature importance
        feature_importance = pd.DataFrame({
            'feature': X.columns,
            'importance': rf.feature_importances_
        }).sort_values('importance', ascending=False)
        
        # Print top features
        print(f"  Top 10 most important features:")
        for i, (feature, importance) in enumerate(zip(feature_importance['feature'].head(10), 
                                                     feature_importance['importance'].head(10))):
            print(f"    {i+1}. {feature}: {importance:.4f}")
        
        if create_visualizations:
            # Plot feature importance
            plt.figure(figsize=(12, 10))
            sns.barplot(x='importance', y='feature', data=feature_importance.head(20))
            plt.title('Feature Importance')
            plt.tight_layout()
            plt.savefig(os.path.join(viz_dir, 'feature_importance.png'))
            plt.close()
        
        # Select features using the model
        # Use the mean importance as threshold
        selector = SelectFromModel(rf, threshold='mean')
        X_selected = selector.fit_transform(X_scaled, y)
        
        print(f"  Original number of features: {X_scaled.shape[1]}")
        print(f"  Number of selected features: {X_selected.shape[1]}")
        
        # Get names of selected features
        selected_feature_mask = selector.get_support()
        selected_features = [feature for feature, selected in zip(feature_names, selected_feature_mask) if selected]
        
        print(f"  Selected features: {', '.join(selected_features[:5])}... and {len(selected_features)-5} more")
        
        # Update X_scaled to only include selected features
        X_scaled = X_selected
        results['statistics']['feature_selection'] = {
            'original_feature_count': len(feature_names),
            'selected_feature_count': len(selected_features),
            'selected_features': selected_features,
            'feature_importance': feature_importance.to_dict()
        }
    else:
        print("\nStep 9: Skipping feature selection")
        selected_features = feature_names
        results['statistics']['feature_selection'] = {
            'performed': False,
            'feature_count': len(feature_names)
        }
    
    # Step 10: Split into training and testing sets
    print("\nStep 10: Splitting into training and testing sets...")
    X_train, X_test, y_train, y_test = train_test_split(
        X_scaled, y, test_size=test_size, random_state=random_state, stratify=y
    )
    
    print(f"  Training set: {X_train.shape[0]} samples")
    print(f"  Testing set: {X_test.shape[0]} samples")
    print(f"  Class distribution in training set: {pd.Series(y_train).value_counts().to_dict()}")
    print(f"  Class distribution in testing set: {pd.Series(y_test).value_counts().to_dict()}")
    
    results['statistics']['train_test_split'] = {
        'train_size': X_train.shape[0],
        'test_size': X_test.shape[0],
        'train_class_distribution': pd.Series(y_train).value_counts().to_dict(),
        'test_class_distribution': pd.Series(y_test).value_counts().to_dict()
    }
    
    # Step 11: Save the preprocessed data
    print("\nStep 11: Saving the preprocessed data...")
    
    # Save the split data
    np.save(os.path.join(output_dir, 'X_train.npy'), X_train)
    np.save(os.path.join(output_dir, 'X_test.npy'), X_test)
    np.save(os.path.join(output_dir, 'y_train.npy'), y_train)
    np.save(os.path.join(output_dir, 'y_test.npy'), y_test)
    
    # Save the scaler for future use
    with open(os.path.join(output_dir, 'scaler.pkl'), 'wb') as f:
        pickle.dump(scaler, f)
    
    # If feature selection was performed, save the selector
    if perform_feature_selection:
        with open(os.path.join(output_dir, 'feature_selector.pkl'), 'wb') as f:
            pickle.dump(selector, f)
    
    # Save the list of selected features
    with open(os.path.join(output_dir, 'selected_features.txt'), 'w') as f:
        for feature in selected_features:
            f.write(f"{feature}\n")
    
    # Save the preprocessed dataset (before split)
    data_to_save = pd.DataFrame(X_scaled, columns=[f"feature_{i}" for i in range(X_scaled.shape[1])])
    data_to_save['is_malware'] = y
    data_to_save.to_csv(os.path.join(output_dir, 'preprocessed_data.csv'), index=False)
    
    # Also save feature mapping
    feature_mapping = {f"feature_{i}": feature for i, feature in enumerate(selected_features)}
    with open(os.path.join(output_dir, 'feature_mapping.json'), 'w') as f:
        import json
        json.dump(feature_mapping, f, indent=2)
    
    # Save preprocessing results
    with open(os.path.join(output_dir, 'preprocessing_results.json'), 'w') as f:
        import json
        # Convert numpy types to Python native types for JSON serialization
        def convert_for_json(obj):
            if isinstance(obj, (np.int64, np.int32, np.int16, np.int8)):
                return int(obj)
            elif isinstance(obj, (np.float64, np.float32, np.float16)):
                return float(obj)
            elif isinstance(obj, np.ndarray):
                return obj.tolist()
            elif isinstance(obj, pd.DataFrame):
                return obj.to_dict()
            else:
                return obj
        
        cleaned_results = {}
        for k, v in results.items():
            if isinstance(v, dict):
                cleaned_results[k] = {k2: convert_for_json(v2) for k2, v2 in v.items()}
            else:
                cleaned_results[k] = convert_for_json(v)
        
        json.dump(cleaned_results, f, indent=2)
    
    print(f"\nPreprocessing completed successfully!")
    print(f"  Preprocessed data saved to {output_dir}")
    print(f"  Visualizations saved to {viz_dir}")
    
    return results

def main():
    """
    Main function to run the preprocessing script
    """
    import argparse
    
    parser = argparse.ArgumentParser(description='Preprocess Android malware dataset')
    parser.add_argument('--input', required=True, help='Input CSV file with cleaned features')
    parser.add_argument('--output', default=None, help='Output directory for preprocessed data')
    parser.add_argument('--test_size', type=float, default=0.2, help='Test set size ratio')
    parser.add_argument('--random_state', type=int, default=42, help='Random seed')
    parser.add_argument('--no_viz', action='store_true', help='Disable visualization generation')
    parser.add_argument('--no_feature_selection', action='store_true', help='Disable feature selection')
    
    args = parser.parse_args()
    
    # If output directory not specified, create one with timestamp
    if not args.output:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        args.output = f"preprocessed_data_{timestamp}"
    
    # Run the preprocessing
    preprocess_android_malware_data(
        input_file=args.input,
        output_dir=args.output,
        test_size=args.test_size,
        random_state=args.random_state,
        create_visualizations=not args.no_viz,
        perform_feature_selection=not args.no_feature_selection
    )

if __name__ == "__main__":
    main()