Nim_detect.py

# $t@$h
# This program does NOT modify your system in any way
# Uses ensemble learning to detect embedded nim files
# Working on ways to improve detection. See comments at
# the bottom for the play by play. Naive bayes + decision tree

# Dataset was built from:
#   Nimlang lexer
#   Github linguist for Nim
#   ReverseTls Tunnel
#   OffensiveNim malware
#   Linux kernel source code
#   Nimskull

import os
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
from sklearn.naive_bayes import MultinomialNB
from sklearn.ensemble import StackingClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score

# Recursively read files in a root directory
def read_files_from_directory(directory):
    contents = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            try:
                with open(os.path.join(root, file), 'r', encoding='utf-8', errors='ignore') as f:
                    content = f.read().replace('\n', ' ')  # Flatten the content into a single line
                    contents.append(content)
            except Exception as e:
                print(f"Error reading {file}: {e}")
    return contents

nim_directory = 'datasets/training/nim'
non_nim_directory = 'datasets/training/non_nim'

nim_samples = read_files_from_directory(nim_directory)
non_nim_samples = read_files_from_directory(non_nim_directory)

# Combining samples and creating labels
samples = nim_samples + non_nim_samples
labels = [1] * len(nim_samples) + [0] * len(non_nim_samples)

# Tailor ngram_range and max_features
vectorizer = TfidfVectorizer(ngram_range=(1, 2), max_features=5000) 
X = vectorizer.fit_transform(samples)

# Split the data
X_train, X_test, y_train, y_test = train_test_split(X, labels, test_size=0.2, random_state=42)

# Define base
base_learners = [
    ('dt', DecisionTreeClassifier(max_depth=15, random_state=42)),
    ('nb', MultinomialNB(alpha=0.1))
]

# Define meta
meta_learner = LogisticRegression(solver='liblinear', random_state=42)

# Stacking classifier
stacked_clf = StackingClassifier(estimators=base_learners, final_estimator=meta_learner)

# Train the stacked model
stacked_clf.fit(X_train, y_train)

# Evaluate accuracy
y_pred = stacked_clf.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)
print(f'Model Accuracy: {accuracy:.2%}')

# Predict with test data
def predict_new_files(model, vectorizer, new_files_directory):
    new_samples = []
    filenames = []
    for root, dirs, files in os.walk(new_files_directory):
        for file in files:
            try:
                with open(os.path.join(root, file), 'r', encoding='utf-8', errors='ignore') as f:
                    new_samples.append(f.read().replace('\n', ' '))
                    filenames.append(file)
            except Exception as e:
                print(f"Error reading {file} in {root}: {e}")

    X_new = vectorizer.transform(new_samples)
    predictions = model.predict(X_new)

    return filenames, predictions

testing_dataset = 'datasets/testing'
filenames, predictions = predict_new_files(stacked_clf, vectorizer, testing_dataset)

# Print filenames and predictions
for filename, pred in zip(filenames, predictions):
    print(f'{filename}: {"Nim" if pred == 1 else "Non-Nim"}')

# Previous results:
# stashakkori$ python3 detect2.py
# Model Accuracy: 98.91%
# reverse_string.txt: Non-Nim
# reverse-string.txt: Non-Nim
# excel_com_bin.txt: Nim
# shellcode_bin.txt: Nim
# palindromic-number.txt: Non-Nim
# connection.txt: Nim
# quick-sort.txt: Non-Nim
# depth-first-search.txt: Non-Nim
# minidump_bin.txt: Nim
# roman-numeral.txt: Non-Nim
# recorder.txt: Nim
# iran_server.txt: Nim
# minimum_spanning_tree.txt: Non-Nim
# main.txt: Nim
# merge-sort.txt: Non-Nim
# pipe.txt: Nim
# dns_resolve.txt: Nim
# roman_numeral.txt: Non-Nim
# dijkstra.txt: Non-Nim
# transpose_matrix.txt: Nim
# longest-common-subsequence.txt: Non-Nim
# quick_sort.txt: Nim
# rot13.txt: Nim
# prime_number.txt: Nim
# prime-number.txt: Non-Nim
# globals.txt: Nim
# selection-sort.txt: Non-Nim
# foreign_server.txt: Nim

# Previous results indicated maybe some overfitting to extension
# E.g some false negatives:
#    foreign_server.nim: Non-Nim (should be Nim)
#    iran_server.nim: Non-Nim (should be Nim)
#    connection.txt: Non-Nim (should be Nim)
# Next steps:
#    1. Add more nim files to nim training data set
#    2. Make all file extensions .txt and retrain

# Current results:
# stashakkori$ python3 detect.py
# Stacked Model Accuracy: 99.47%
# dijkstra.cpp: Non-Nim
# prime_number.py: Non-Nim
# merge-sort.cpp: Non-Nim
# shellcode_bin.nim: Nim
# quick_sort.py: Non-Nim
# longest-common-subsequence.cpp: Non-Nim
# reverse_string.py: Non-Nim
# excel_com_bin.nim: Nim
# minidump_bin.nim: Nim
# connection.txt: Non-Nim
# prime-number.cpp: Non-Nim
# minimum_spanning_tree.py: Non-Nim
# roman_numeral.py: Non-Nim
# rot13.py: Non-Nim
# selection-sort.cpp: Non-Nim
# main.nim: Non-Nim
# iran_server.nim: Non-Nim
# recorder.nim: Nim
# reverse-string.cpp: Non-Nim
# dns_resolve.nim: Nim
# pipe.nim: Nim
# transpose_matrix.py: Non-Nim
# quick-sort.cpp: Non-Nim
# rot13.txt: Non-Nim
# depth-first-search.cpp: Non-Nim
# palindromic-number.cpp: Non-Nim
# foreign_server.nim: Non-Nim
# globals.txt: Non-Nim
# roman-numeral.cpp: Non-Nim

# Implementing next steps and adding the changes below better detects testing data.
# Happy with where this is now, others can further this work.
#   a. Flattened files - replace newlines with spaces
#   b. Changed vectorization to n-gram range instead of unigram
#   c. Dimensionality reduction to reduce noise as well. Seems to have worked.
#   d. Increased max depth of decision tree but not too much to overfit
#   e. Adjusted Laplace smoothing of Bayes because suspected overfitting
#   f. Set solver to liblinear. Did this for performance but maybe improved accuracy.