import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import confusion_matrix, accuracy_score
import pickle
import os
import subprocess
import sys
import csv 

# Load dataset (dummy dataset provided as an example)
def load_dataset():
    # Sample data for demonstration purposes. Replace this with actual dataset.
    data = pd.read_csv('results.csv')
    return data
def md5_hash(file_path):
    """Compute the MD5 hash of a file using md5sum."""
    result = subprocess.run(['md5sum', file_path], capture_output=True, text=True)
    return result.stdout.split()[0]


def get_machine_type(file_path):
    """Get the machine architecture from an ELF file using readelf."""
    try:
        result = subprocess.run(['readelf', '-h', file_path], capture_output=True, text=True)
        for line in result.stdout.splitlines():
            if 'Machine:' in line:
                return line.split(':')[1].strip()
    except Exception as e:
        print(f"Error getting machine type: {e}")
    return None


def get_number_of_sections(file_path):
    """Get the number of sections in an ELF file using readelf."""
    try:
        result = subprocess.run(['readelf', '-h', file_path], capture_output=True, text=True)
        for line in result.stdout.splitlines():
            if 'Number of section headers:' in line:
                return int(line.split(':')[1].strip())
    except Exception as e:
        print(f"Error getting number of sections: {e}")
    return None


def get_resource_size(file_path):
    """Get the size of the .rodata section (resources) in an ELF file using readelf."""
    try:
        result = subprocess.run(['readelf', '-S', file_path], capture_output=True, text=True)
        for line in result.stdout.splitlines():
            if '.rodata' in line:
                size_hex = line.split()[5]
                return int(size_hex, 16)  # Convert from hex to decimal
    except Exception as e:
        print(f"Error getting resource size: {e}")
    return 0


def get_linker_version(file_path):
    """Get the linker version from an ELF file using objdump."""
    try:
        result = subprocess.run(['objdump', '-p', file_path], capture_output=True, text=True)
        for line in result.stdout.splitlines():
            if 'Version:' in line:
                version = line.split(':')[1].strip()
                major_version = version.split('.')[0]
                minor_version = version.split('.')[1] if '.' in version else '0'
                return major_version, minor_version
    except Exception as e:
        print(f"Error getting linker version: {e}")
    return None, None


def get_dynamic_info(file_path):
    """Get dynamic linking information (e.g., import address table equivalent) using readelf."""
    try:
        result = subprocess.run(['readelf', '-d', file_path], capture_output=True, text=True)
        dynamic_info = []
        for line in result.stdout.splitlines():
            dynamic_info.append(line)
        return dynamic_info
    except Exception as e:
        print(f"Error getting dynamic linking info: {e}")
    return None


def extract_features(file_path):
    """Extract features from an ELF file."""
    features = {
        'FileName': file_path,
        'md5Hash': md5_hash(file_path),
        'Machine': get_machine_type(file_path),
        'NumberOfSections': get_number_of_sections(file_path),
        'ResourceSize': get_resource_size(file_path),
        'LinkerVersionMajor': 0,
        'LinkerVersionMinor': 0,
        'DynamicInfo': get_dynamic_info(file_path)
    }

    # Get linker version
    major_version, minor_version = get_linker_version(file_path)
    features['LinkerVersionMajor'] = major_version
    features['LinkerVersionMinor'] = minor_version

  

    return features

# Train the model
def train_model(data):
    # Split into features and labels
    # X = data.drop(columns=['RansomwareType','FileName', 'md5Hash', 'Machine','DynamicInfo'])  # Features

    X = data.drop(columns=['RansomwareType','FileName', 'md5Hash', 'Machine'])  # Features
    y = data['RansomwareType']                # Labels (target)

    # Split the dataset into training and testing sets
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # Train the RandomForestClassifier
    clf = RandomForestClassifier()
    clf.fit(X_train, y_train)

    # Save the model
    with open('model.pkl', 'wb') as f:
        pickle.dump(clf, f)

    # Evaluate the model on the test set
    y_pred = clf.predict(X_test)
    print("Confusion Matrix:")
    print(confusion_matrix(y_test, y_pred))
    print("Accuracy:", accuracy_score(y_test, y_pred))


if __name__ == "__main__":
    # Load dataset and train the model
    data = load_dataset()
    train_model(data)

