PreExtractFeatures.py

import numpy as np
import os
import argparse
import pefile
from pathlib import Path
import random
from tqdm import tqdm
import multiprocessing
import re

# --- Configuration ---
NUM_FEATURES = 11

# --- Dataset Paths ---
TRAIN_ROOT = "data_filtered/32bit/train"
VAL_ROOT = "data_filtered/32bit/val"
CHALLENGE_ROOT = "data_separated/32bit/challenge"
BENIGN_PACKED_DIR = "data_filtered/32bit/train/goodware"
MALWARE_TRAIN_DIR = "data_filtered/32bit/train/malware"

# --- Output Directory ---
OUTPUT_DIR = "data_separated_features32bit"

# --- Suspicious Pattern Definitions ---
SUSPICIOUS_IMPORTS_NETWORKING = [
    re.compile(rb'Socket', re.IGNORECASE),
    re.compile(rb'InternetOpen', re.IGNORECASE),
    re.compile(rb'InternetReadFile', re.IGNORECASE),
    re.compile(rb'HttpSendRequest', re.IGNORECASE),
    re.compile(rb'WSAStartup', re.IGNORECASE),
    re.compile(rb'connect', re.IGNORECASE),
    re.compile(rb'URLDownloadToFile', re.IGNORECASE),
]

SUSPICIOUS_IMPORTS_ANTIDEBUG = [
    re.compile(rb'IsDebuggerPresent', re.IGNORECASE),
    re.compile(rb'CheckRemoteDebuggerPresent', re.IGNORECASE),
    re.compile(rb'NtQueryInformationProcess', re.IGNORECASE),
    re.compile(rb'OutputDebugString', re.IGNORECASE),
]


def calculate_entropy(data):
    """Calculate Shannon entropy of bytes"""
    if len(data) == 0:
        return 0.0
    
    byte_counts = np.bincount(np.frombuffer(data, dtype=np.uint8), minlength=256)
    probabilities = byte_counts / len(data)
    entropy = -np.sum(probabilities[probabilities > 0] * np.log2(probabilities[probabilities > 0]))
    return float(entropy)


def get_features(pe, file_bytes):
    
    features = {}
    
    # 1. has_version_info
    features['has_version_info'] = 0.0
    if hasattr(pe, 'DIRECTORY_ENTRY_RESOURCE'):
        try:
            for entry in pe.DIRECTORY_ENTRY_RESOURCE.entries:
                if entry.id == pefile.RESOURCE_TYPE['RT_VERSION']:
                    features['has_version_info'] = 1.0
                    break
        except Exception:
            pass
    
    # 2. has_signature
    features['has_signature'] = 0.0
    if hasattr(pe, 'OPTIONAL_HEADER'):
        sec_dir_index = pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_SECURITY']
        if len(pe.OPTIONAL_HEADER.DATA_DIRECTORY) > sec_dir_index:
            dir_entry = pe.OPTIONAL_HEADER.DATA_DIRECTORY[sec_dir_index]
            if dir_entry.Size > 0 and dir_entry.VirtualAddress > 0:
                features['has_signature'] = 1.0
    
    # 3. has_debug_info
    features['has_debug_info'] = 0.0
    debug_dir_index = pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_DEBUG']
    if hasattr(pe, 'OPTIONAL_HEADER') and len(pe.OPTIONAL_HEADER.DATA_DIRECTORY) > debug_dir_index:
        if pe.OPTIONAL_HEADER.DATA_DIRECTORY[debug_dir_index].Size > 0:
            features['has_debug_info'] = 1.0
    
    # 4. imports_networking_count
    features['imports_networking_count'] = 0.0
    if hasattr(pe, 'DIRECTORY_ENTRY_IMPORT'):
        for entry in pe.DIRECTORY_ENTRY_IMPORT:
            for imp in entry.imports:
                if imp.name:
                    for susp_regex in SUSPICIOUS_IMPORTS_NETWORKING:
                        if susp_regex.search(imp.name):
                            features['imports_networking_count'] += 1.0
                            break
    
    # 5. imports_antidebug_count
    features['imports_antidebug_count'] = 0.0
    if hasattr(pe, 'DIRECTORY_ENTRY_IMPORT'):
        for entry in pe.DIRECTORY_ENTRY_IMPORT:
            for imp in entry.imports:
                if imp.name:
                    for susp_regex in SUSPICIOUS_IMPORTS_ANTIDEBUG:
                        if susp_regex.search(imp.name):
                            features['imports_antidebug_count'] += 1.0
                            break
    
    # 6. dll_char_dynamic_base
    features['dll_char_dynamic_base'] = 0.0
    if hasattr(pe, 'OPTIONAL_HEADER'):
        if pe.OPTIONAL_HEADER.DllCharacteristics & 0x0040:
            features['dll_char_dynamic_base'] = 1.0
    
    # 7. dll_char_nx_compat
    features['dll_char_nx_compat'] = 0.0
    if hasattr(pe, 'OPTIONAL_HEADER'):
        if pe.OPTIONAL_HEADER.DllCharacteristics & 0x0100:
            features['dll_char_nx_compat'] = 1.0
    
    # 8. is_32bit
    features['is_32bit'] = 0.0
    if hasattr(pe, 'FILE_HEADER'):
        if pe.FILE_HEADER.Machine == pefile.MACHINE_TYPE['IMAGE_FILE_MACHINE_I386']:
            features['is_32bit'] = 1.0
    
    # 9. opt_magic
    features['opt_magic'] = 0.0
    if hasattr(pe, 'OPTIONAL_HEADER'):
        features['opt_magic'] = float(pe.OPTIONAL_HEADER.Magic)
    
    
    # 10. section_entropy_variance
    features['section_entropy_variance'] = 0.0
    if hasattr(pe, 'sections'):
        section_entropies = []
        for section in pe.sections:
            try:
                if section.SizeOfRawData > 0:
                    section_data = pe.get_data(section.VirtualAddress, min(section.SizeOfRawData, 8192))
                    entropy = calculate_entropy(section_data)
                    section_entropies.append(entropy)
            except:
                pass
        
        if len(section_entropies) > 1:
            features['section_entropy_variance'] = float(np.var(section_entropies))
    
    # 11. code_to_data_ratio 
    features['code_to_data_ratio'] = 0.0
    if hasattr(pe, 'sections'):
        code_sections = sum(1 for s in pe.sections if s.Characteristics & 0x00000020)
        data_sections = sum(1 for s in pe.sections if s.Characteristics & 0x00000040)
        if data_sections > 0:
            features['code_to_data_ratio'] = float(code_sections) / float(data_sections)
    
    
    # --- FINAL FEATURE VECTOR ---
    feature_order = [
        'has_version_info',
        'has_signature',
        'has_debug_info',
        'imports_networking_count',
        'imports_antidebug_count',
        'dll_char_dynamic_base',
        'dll_char_nx_compat',
        'is_32bit',
        'opt_magic',
        'section_entropy_variance',
        'code_to_data_ratio'
    ]
    
    if len(feature_order) != NUM_FEATURES:
        raise Exception(f"Feature mismatch! NUM_FEATURES is {NUM_FEATURES} but feature_order has {len(feature_order)}")
    
    return np.array([features[f] for f in feature_order], dtype=np.float32)


def process_file(file_path_label):
    """Worker function to process a single file"""
    file_path, label = file_path_label
    
    try:
        file_bytes = Path(file_path).read_bytes()
        pe = pefile.PE(data=file_bytes, fast_load=False)
        feature_vector = get_features(pe, file_bytes)
    except Exception as e:
        return None
    finally:
        if 'pe' in locals() and locals()['pe']:
            try:
                locals()['pe'].close()
            except:
                pass
    
    return feature_vector, label


def get_file_paths_and_labels(packed_dir, malware_dir, split_name):
    """Collect file paths and labels for a dataset split"""
    
    def find_files(directory):
        paths = []
        if not os.path.isdir(directory):
            print(f"Warning: Directory not found, skipping: {directory}")
            return paths
        for root, _, filenames in os.walk(directory):
            for filename in filenames:
                paths.append(Path(os.path.join(root, filename)))
        return paths

    goodware_paths = find_files(packed_dir)
    malware_paths = find_files(malware_dir)

    print(f"[{split_name.upper()}] Found {len(goodware_paths)} benign files and {len(malware_paths)} malware files.")
    
    all_files = [(str(f), 0) for f in goodware_paths] + [(str(f), 1) for f in malware_paths]
    if not all_files:
        return []
    
    random.shuffle(all_files)
    return all_files


def main(num_workers):
    """Main preprocessing pipeline"""
    random.seed(42)
    np.random.seed(42)

    os.makedirs(OUTPUT_DIR, exist_ok=True)

    # Collect file paths for all splits
    train_files = get_file_paths_and_labels(BENIGN_PACKED_DIR, MALWARE_TRAIN_DIR, 'train')
    val_files = get_file_paths_and_labels(
        Path(VAL_ROOT).joinpath('goodware'),
        Path(VAL_ROOT).joinpath('malware'),
        'val'
    )
    challenge_files = get_file_paths_and_labels(
        Path(CHALLENGE_ROOT).joinpath('goodware'),
        Path(CHALLENGE_ROOT).joinpath('malware'),
        'challenge'
    )

    data_splits = {
        'train': train_files,
        'val': val_files,
        'challenge': challenge_files
    }
    
    feature_names = [
        'has_version_info',
        'has_signature',
        'has_debug_info',
        'imports_networking_count',
        'imports_antidebug_count',
        'dll_char_dynamic_base',
        'dll_char_nx_compat',
        'is_32bit',
        'opt_magic',
        'section_entropy_variance',
        'code_to_data_ratio'
    ]
    
    assert len(feature_names) == NUM_FEATURES, \
        f"Mismatch: NUM_FEATURES is {NUM_FEATURES} but feature_names has {len(feature_names)} items."

    feature_names_path = os.path.join(OUTPUT_DIR, 'feature_names.txt')
    with open(feature_names_path, 'w') as f:
        f.write(f"HYBRID FEATURE SET - Original 9 + 4 New ({len(feature_names)} features)\n")
        f.write("=" * 80 + "\n")
        f.write("TIER 0 - Original 9 (Proven Generalization):\n")
        for i, name in enumerate(feature_names[0:9], 1):
            f.write(f"  {i}. {name}\n")
        f.write("\nTIER 1 - New Non-Redundant Features (Orthogonal Signals):\n")
        for i, name in enumerate(feature_names[9:13], 10):
            f.write(f"  {i}. {name}\n")
        f.write("\nRATIONALE:\n")
        f.write("- Kept all original 9 for stability and proven generalization\n")
        f.write("- Added 4 orthogonal features: section analysis, structural anomalies\n")
        f.write("- Excluded engineered/redundant features that caused overfitting\n")
        f.write("- Expected improvement over original 9, much better than failed 15-feature set\n")
    print(f"Saved hybrid feature set ({len(feature_names)} features) to {feature_names_path}")
    
    ctx = multiprocessing.get_context("spawn")
    
    for split_name, files in data_splits.items():
        if not files:
            print(f"Warning: No files found for {split_name}. Skipping.")
            continue

        print(f"\n--- Processing {split_name.upper()} Data ({len(files)} files) ---")
        
        X_list = []
        y_list = []
        filenames_list = []
        
        with ctx.Pool(processes=num_workers) as pool:
            results_generator = pool.imap(process_file, files, chunksize=256)

            paired_iterator = zip(files, results_generator)

            for (original_file_tuple, result) in tqdm(paired_iterator, total=len(files), desc=f"{split_name} conversion"):
                
                if result is None:
                    continue
                
                feat, label = result
                X_list.append(feat)
                y_list.append(label)
                
                original_filepath = original_file_tuple[0]
                filenames_list.append(os.path.basename(original_filepath))
        
        if not X_list:
            print(f"Warning: No features were extracted for {split_name}. Skipping save.")
            continue

        X_data = np.array(X_list, dtype=np.float32)
        y_data = np.array(y_list, dtype=np.int32)
        
        X_path = os.path.join(OUTPUT_DIR, f"X_{split_name}.npy")
        y_path = os.path.join(OUTPUT_DIR, f"y_{split_name}.npy")
        
        np.save(X_path, X_data)
        np.save(y_path, y_data)

        print(f"Saved {split_name} data: X shape {X_data.shape}, y shape {y_data.shape}")
        
        if split_name == 'challenge':
            filenames_path = os.path.join(OUTPUT_DIR, 'challenge_filenames.txt')
            with open(filenames_path, 'w') as f:
                for filename in filenames_list:
                    f.write(f"{filename}\n")
            print(f"Saved {len(filenames_list)} challenge filenames to {filenames_path}")
        
        print(f"\n{split_name.upper()} Feature Statistics:")
        for i in range(min(10, NUM_FEATURES)):
            mean_val = X_data.mean(axis=0)[i]
            std_val = X_data.std(axis=0)[i]
            print(f"  {i+1:2d}. {feature_names[i]:40s}: mean={mean_val:8.4f}, std={std_val:8.4f}")

    print("\n--- Preprocessing Complete ---")
    print(f"Total features extracted: {NUM_FEATURES}")
    print(f"Output directory: {OUTPUT_DIR}")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Extract feature set'
    )
    parser.add_argument(
        '--num_workers',
        type=int,
        default=8,
        help='Number of parallel processes to use.'
    )
    args = parser.parse_args()
    main(args.num_workers)