Simple decision tree implemented in scikit-learn, trained on dummy malware data

malware_data.csv

MalwareTree.py

"""
A simple decision tree implementation in scikit-learn.
Trained on fabricated malware data.

Features: file id, name, size in MB, extension and executability.
Labels: malware or not.

Output:

    Accuracy:  0.75
    Predictions:  [1 1 1 1 1 1 0 0 0 0 0 0]

"""

#  ---------------- Importing Libraries ----------------
import pandas as pd
from sklearn.tree import DecisionTreeClassifier
from sklearn.model_selection import train_test_split
#  -----------------------------------------------------

# Load the data
data = pd.read_csv('malware_data.csv')

# Split the data into features and labels
X = data.drop(columns=['label', 'file_id', 'file_name'])
y = data['label']

# convert to numerical data
X = pd.get_dummies(X)
y = y.astype('int')

# Split the data 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 decision tree model
model = DecisionTreeClassifier()
model.fit(X_train, y_train)

# Test the model on the testing data
accuracy = model.score(X_test, y_test)
print('Accuracy: ', accuracy)

# Predict on new, unseen data
new_data = pd.read_csv('new_malware_data.csv')
new_data = new_data.drop(columns=['file_id', 'file_name'])
new_data = pd.get_dummies(new_data)
predictions = model.predict(new_data)
print('Predictions: ', predictions)

new_malware_data.csv