Malware Detection System

# ── Data handling ──────────────────────────────────────────────────────────
import numpy as np          # numerical arrays and fast math operations
import pandas as pd         # load and manipulate the CSV dataset as a DataFrame

# ── Visualisation ──────────────────────────────────────────────────────────
import seaborn as sns        # statistical charts: heatmap, boxplot, histogram
import matplotlib.pyplot as plt  # core plotting engine used by seaborn

# ── Model evaluation ───────────────────────────────────────────────────────
from sklearn.metrics import (
    accuracy_score,           # proportion of correct predictions: (TP+TN)/Total
    classification_report,   # precision, recall, F1-score, support per class
    confusion_matrix          # 2×2 table of TP, TN, FP, FN
)

# ── Data splitting ─────────────────────────────────────────────────────────
from sklearn.model_selection import train_test_split  # 80/20 split

# ── Preprocessing ──────────────────────────────────────────────────────────
from sklearn.preprocessing import MinMaxScaler  # scale features to [0, 1]

# ── ML Algorithms ──────────────────────────────────────────────────────────
from sklearn.ensemble       import RandomForestClassifier  # best model
from sklearn.neighbors      import KNeighborsClassifier     # distance-based
from sklearn.linear_model   import LogisticRegression       # linear baseline
from sklearn.tree           import DecisionTreeClassifier    # single tree

# ── Model persistence ──────────────────────────────────────────────────────
import pickle  # save and load trained models to/from disk files (.pkl)

print("✅ All libraries imported successfully")

# Load the CSV file into a pandas DataFrame
# pd.read_csv() reads the file line by line and creates a table in memory
# where each row = one file, each column = one extracted PE feature
data = pd.read_csv("dataset_malwares.csv")

# ── Quick shape check ──────────────────────────────────────────────────────
# .shape returns a tuple: (number_of_rows, number_of_columns)
# number_of_rows   = how many executable files are in the dataset
# number_of_columns = how many features + the target label column
print(f"Dataset shape: {data.shape}")
print(f"  → {data.shape[0]:,} files (rows)")
print(f"  → {data.shape[1]} columns (features + target)")

# .head() displays the first 5 rows of the DataFrame
# This lets you see what the raw data looks like before any processing
# Look for: column names, data types at a glance, and the Malware target column
data.head()

# .info() is the most important first inspection call:
# • Lists every column name
# • Shows the data type of each column (int64, float64, object)
# • Shows how many non-null (non-missing) values each column has
# • Reports total memory usage
# If any column shows fewer non-null entries than the total row count,
# those missing values need to be handled before training.
data.info()

print("\n──── Missing value check ────")
print(data.isnull().sum())   # print count of NaN values per column

# Check how many malware vs benign samples we have
# Class imbalance — e.g. 95% benign, 5% malware — would cause the model
# to "cheat" by always predicting benign and still appear 95% accurate
class_counts = data['Malware'].value_counts()
class_counts.index = class_counts.index.astype(int)
print("Class distribution:")
benign_count  = class_counts.get(0, 0)
malware_count = class_counts.get(1, 0)
total = len(data)
print(f"  Benign  (0): {benign_count:,} files")
print(f"  Malware (1): {malware_count:,} files")
print(f"  Balance: {benign_count/total*100:.1f}% benign / {malware_count/total*100:.1f}% malware")

# Visual representation with a simple bar chart
plt.figure(figsize=(6, 4))
sns.countplot(x='Malware', data=data,
             hue='Malware',
             palette={0: '#00e676', 1: '#ff4d6d'},
             legend=False)
plt.title('Class Distribution: Benign vs Malware', fontweight='bold')
plt.xticks([0, 1], ['Benign (0)', 'Malware (1)'])
plt.ylabel('Number of Files')
plt.show()

python --version

winget install Python.Python.3

pip install pefile tqdm pandas

# ============================================================
#  collect_benign.py  —  run on your Windows PC, NOT in Colab
#  Scans System32, SysWOW64, Program Files for .exe/.dll files,
#  extracts PE features, combines with dataset_malwares.csv,
#  and saves combined_dataset.csv ready to upload to Colab.
# ============================================================

import os, glob, math, random
import pandas as pd
import pefile
from tqdm import tqdm

# ── Configuration ─────────────────────────────────────────────────────────
BENIGN_SOURCES = [
    r"C:\Windows\System32",
    r"C:\Windows\SysWOW64",
    r"C:\Program Files",
    r"C:\Program Files (x86)",
]
MAX_PER_DIR = 2000                        # raise for a larger dataset
MALWARE_CSV = "dataset_malwares.csv"       # must be in the same folder as this script
OUTPUT_CSV  = "combined_dataset.csv"       # this is what you upload to Colab

# ── Shannon entropy helper ─────────────────────────────────────────────────
def safe_entropy(data):
    if not data: return 0.0
    freq = [0] * 256
    for b in data: freq[b] += 1
    n = len(data)
    return -sum(c/n * math.log2(c/n) for c in freq if c)

# ── PE feature extractor ───────────────────────────────────────────────────
def extract_pe_features(filepath):
    # Returns a feature dict matching dataset_malwares.csv columns,
    # or None if the file is not a valid PE / cannot be read.
    try:
        pe = pefile.PE(filepath, fast_load=False)
    except:
        return None
    try:
        secs = pe.sections if pe.sections else []
        ents, raw_s, virt_s, phys, virts, ptrs, chars = [], [], [], [], [], [], []
        total_len = 0
        for s in secs:
            ents.append(safe_entropy(s.get_data()))
            total_len += s.SizeOfRawData
            raw_s.append(s.SizeOfRawData);  virt_s.append(s.Misc_VirtualSize)
            phys.append(s.SizeOfRawData);   virts.append(s.VirtualAddress)
            ptrs.append(s.PointerToRawData); chars.append(s.Characteristics)

        opt, fh = pe.OPTIONAL_HEADER, pe.FILE_HEADER
        dd = opt.DATA_DIRECTORY if hasattr(opt, 'DATA_DIRECTORY') else []

        SUSP_API = {'VirtualAlloc', 'VirtualProtect', 'WriteProcessMemory',
                    'CreateRemoteThread', 'OpenProcess', 'SetWindowsHookEx',
                    'RegSetValueEx', 'CreateMutex', 'InternetOpen', 'URLDownloadToFile'}
        NORM_SEC = {b'.text', b'.data', b'.rdata', b'.rsrc',
                    b'.reloc', b'.bss', b'.idata', b'.edata', b'.pdata'}

        susp_imp = dir_imp = dir_imp_sz = dir_exp = 0
        if hasattr(pe, 'DIRECTORY_ENTRY_IMPORT'):
            dir_imp = len(pe.DIRECTORY_ENTRY_IMPORT)
            for e in pe.DIRECTORY_ENTRY_IMPORT:
                dir_imp_sz += len(e.imports)
                for i in e.imports:
                    if i.name and i.name.decode(errors='ignore') in SUSP_API:
                        susp_imp += 1
        if hasattr(pe, 'DIRECTORY_ENTRY_EXPORT'):
            dir_exp = len(pe.DIRECTORY_ENTRY_EXPORT.symbols)
        susp_sec = sum(1 for s in secs if s.Name.rstrip(b'') not in NORM_SEC)

        features = {
            'Name': os.path.basename(filepath),
            'e_magic': pe.DOS_HEADER.e_magic,   'e_cblp': pe.DOS_HEADER.e_cblp,
            'e_cp': pe.DOS_HEADER.e_cp,           'e_crlc': pe.DOS_HEADER.e_crlc,
            'e_cparhdr': pe.DOS_HEADER.e_cparhdr, 'e_minalloc': pe.DOS_HEADER.e_minalloc,
            'e_maxalloc': pe.DOS_HEADER.e_maxalloc,'e_ss': pe.DOS_HEADER.e_ss,
            'e_sp': pe.DOS_HEADER.e_sp,           'e_csum': pe.DOS_HEADER.e_csum,
            'e_ip': pe.DOS_HEADER.e_ip,           'e_cs': pe.DOS_HEADER.e_cs,
            'e_lfarlc': pe.DOS_HEADER.e_lfarlc,  'e_ovno': pe.DOS_HEADER.e_ovno,
            'e_oemid': pe.DOS_HEADER.e_oemid,    'e_oeminfo': pe.DOS_HEADER.e_oeminfo,
            'e_lfanew': pe.DOS_HEADER.e_lfanew,
            'Machine': fh.Machine,                'NumberOfSections': fh.NumberOfSections,
            'TimeDateStamp': fh.TimeDateStamp,    'PointerToSymbolTable': fh.PointerToSymbolTable,
            'NumberOfSymbols': fh.NumberOfSymbols,'SizeOfOptionalHeader': fh.SizeOfOptionalHeader,
            'Characteristics': fh.Characteristics,'Magic': opt.Magic,
            'MajorLinkerVersion': opt.MajorLinkerVersion,
            'MinorLinkerVersion': opt.MinorLinkerVersion,
            'SizeOfCode': opt.SizeOfCode,
            'SizeOfInitializedData': opt.SizeOfInitializedData,
            'SizeOfUninitializedData': opt.SizeOfUninitializedData,
            'AddressOfEntryPoint': opt.AddressOfEntryPoint,
            'BaseOfCode': opt.BaseOfCode,        'ImageBase': opt.ImageBase,
            'SectionAlignment': opt.SectionAlignment,
            'FileAlignment': opt.FileAlignment,
            'MajorOperatingSystemVersion': opt.MajorOperatingSystemVersion,
            'MinorOperatingSystemVersion': opt.MinorOperatingSystemVersion,
            'MajorImageVersion': opt.MajorImageVersion,
            'MinorImageVersion': opt.MinorImageVersion,
            'MajorSubsystemVersion': opt.MajorSubsystemVersion,
            'MinorSubsystemVersion': opt.MinorSubsystemVersion,
            'SizeOfImage': opt.SizeOfImage,      'SizeOfHeaders': opt.SizeOfHeaders,
            'CheckSum': opt.CheckSum,            'Subsystem': opt.Subsystem,
            'DllCharacteristics': opt.DllCharacteristics,
            'SizeOfStackReserve': opt.SizeOfStackReserve,
            'SizeOfStackCommit': opt.SizeOfStackCommit,
            'SizeOfHeapReserve': opt.SizeOfHeapReserve,
            'SizeOfHeapCommit': opt.SizeOfHeapCommit,
            'LoaderFlags': opt.LoaderFlags,
            'NumberOfRvaAndSizes': opt.NumberOfRvaAndSizes,
            'Malware': 0,   # 0 = benign — all files here are known-good
            'SuspiciousImportFunctions': susp_imp,
            'SuspiciousNameSection': susp_sec,
            'SectionsLength': total_len,
            'SectionMinEntropy':     min(ents)  if ents  else 0,
            'SectionMaxEntropy':     max(ents)  if ents  else 0,
            'SectionMinRawsize':     min(raw_s) if raw_s else 0,
            'SectionMaxRawsize':     max(raw_s) if raw_s else 0,
            'SectionMinVirtualsize': min(virt_s)if virt_selse 0,
            'SectionMaxVirtualsize': max(virt_s)if virt_selse 0,
            'SectionMaxPhysical':    max(phys)  if phys  else 0,
            'SectionMinPhysical':    min(phys)  if phys  else 0,
            'SectionMaxVirtual':     max(virts) if virts else 0,
            'SectionMinVirtual':     min(virts) if virts else 0,
            'SectionMaxPointerData': max(ptrs)  if ptrs  else 0,
            'SectionMinPointerData': min(ptrs)  if ptrs  else 0,
            'SectionMaxChar':        max(chars) if chars else 0,
            'SectionMainChar':       secs[0].Characteristics if secs else 0,
            'DirectoryEntryImport':   dir_imp,
            'DirectoryEntryImportSize': dir_imp_sz,
            'DirectoryEntryExport':   dir_exp,
            'ImageDirectoryEntryExport':    dd[0].Size if len(dd) > 0 else 0,
            'ImageDirectoryEntryImport':    dd[1].Size if len(dd) > 1 else 0,
            'ImageDirectoryEntryResource':  dd[2].Size if len(dd) > 2 else 0,
            'ImageDirectoryEntryException': dd[3].Size if len(dd) > 3 else 0,
            'ImageDirectoryEntrySecurity':  dd[4].Size if len(dd) > 4 else 0,
        }
        pe.close()
        return features
    except:
        try: pe.close()
        except: pass
        return None

# ── Directory scanner ──────────────────────────────────────────────────────
def collect_benign_pe_files(source_dirs, max_per_dir):
    all_records = []
    for source_dir in source_dirs:
        if not os.path.isdir(source_dir):
            print(f"⚠️  Skipping (not found): {source_dir}"); continue
        pe_files = []
        for ext in ("*.exe", "*.dll"):
            pe_files.extend(
                glob.glob(os.path.join(source_dir, "**", ext), recursive=True)
            )
        random.shuffle(pe_files)
        pe_files = pe_files[:max_per_dir]
        print(f"
📂 Scanning: {source_dir}  ({len(pe_files)} files)")
        ok = skipped = 0
        for path in tqdm(pe_files, unit="file"):
            rec = extract_pe_features(path)
            if rec: all_records.append(rec); ok += 1
            else: skipped += 1
        print(f"   ✅ Parsed: {ok}  |  ⚠️  Skipped: {skipped}")
    print(f"
✅ Total benign records collected: {len(all_records):,}")
    return all_records

# ── Combine with original malware CSV ─────────────────────────────────────
def combine_with_malware(benign_records, malware_csv_path):
    print(f"
📄 Loading: {malware_csv_path}")
    malware_df  = pd.read_csv(malware_csv_path)
    benign_df   = pd.DataFrame(benign_records)
    cols        = malware_df.columns.tolist()
    # Align columns: fill any missing columns with 0, drop any extras
    combined    = pd.concat(
        [malware_df, benign_df.reindex(columns=cols, fill_value=0)],
        ignore_index=True
    )
    before      = len(combined)
    combined.drop_duplicates(
        subset=[c for c in cols if c != 'Name'], keep='first', inplace=True
    )
    combined = combined.sample(frac=1, random_state=42).reset_index(drop=True)
    total = len(combined)
    b = (combined['Malware'] == 0).sum()
    m = (combined['Malware'] == 1).sum()
    print(f"
{'='*45}")
    print(f"  Total rows     : {total:,}")
    print(f"  Benign  (0)    : {b:,}  ({b/total*100:.1f}%)")
    print(f"  Malware (1)    : {m:,}  ({m/total*100:.1f}%)")
    print(f"  Dupes removed  : {before - total}")
    print(f"{'='*45}")
    return combined

# ── Main ───────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    if not os.path.isfile(MALWARE_CSV):
        print(f"❌ '{MALWARE_CSV}' not found.")
        print(f"   Place dataset_malwares.csv in the same folder as this script.")
        exit(1)
    print("="*45 + "
  Benign PE Feature Collector
" + "="*45)
    records  = collect_benign_pe_files(BENIGN_SOURCES, MAX_PER_DIR)
    if not records:
        print("❌ No records collected. Check directory paths."); exit(1)
    combined = combine_with_malware(records, MALWARE_CSV)
    combined.to_csv(OUTPUT_CSV, index=False)
    print(f"
✅ Saved → {OUTPUT_CSV}")
    print("   Upload this file to Colab and update pd.read_csv() in Section 3.")
    print("
Done! 🎉")

python collect_benign.py

# Change FROM this:
data = pd.read_csv("dataset_malwares.csv")

# Change TO this:
data = pd.read_csv("combined_dataset.csv")

# Overall statistical summary of all numerical features
print("── Full dataset statistics ──")
print(data.describe().round(2))

# Compare means between malware (1) and benign (0) for each feature
# Large differences between the two groups = high predictive value
print("\n── Mean feature values by class ──")
print(data.groupby('Malware').mean(numeric_only=True).round(3).T)  # .T transposes for readability

# KDE (Kernel Density Estimate) plot: a smooth histogram showing the
# distribution of entropy values for each class overlaid on the same chart
# We expect malware (red) to peak at higher entropy values than benign (green)

fig, axes = plt.subplots(1, 3, figsize=(16, 5))
entropy_cols = ['SectionMaxEntropy', 'SectionMinEntropy', 'SectionsLength']

for ax, col in zip(axes, entropy_cols):
    for label, colour in [(0, '#00e676'), (1, '#ff4d6d')]:
        subset = data[data['Malware'] == label][col].dropna()
        if len(subset) > 1:
            ax.hist(subset, bins=50, alpha=0.55, color=colour,
                   density=True, label='Benign' if label==0 else 'Malware')
    ax.set_title(col, fontweight='bold')
    ax.set_xlabel('Entropy (bits)')
    ax.set_ylabel('Density')
    ax.legend()
    ax.axvline(7.0, color='#ffb300', linestyle='--', linewidth=1,
               label='Entropy = 7.0 threshold')

plt.suptitle('Section Entropy Distribution: Malware vs Benign',
            fontsize=14, fontweight='bold', y=1.02)
plt.tight_layout()
plt.show()

# Compute the correlation matrix for all numerical columns
# .corr() computes Pearson correlation coefficient between every column pair

# Drop the non-numeric columns we will remove later anyway
numeric_data = data.drop(['Name', 'Machine', 'TimeDateStamp', 'SectionMainChar'], axis=1)

corr_matrix = numeric_data.corr()

plt.figure(figsize=(14, 11))
sns.heatmap(
    corr_matrix,
    annot=True,           # write the correlation number in each cell
    fmt='.2f',            # format to 2 decimal places
    cmap='RdYlGn',       # red = negative, yellow = near zero, green = positive
    center=0,             # centre the colour scale at 0
    linewidths=0.3,
    linecolor='#1c2436'
)
plt.title('Feature Correlation Matrix — Malware Dataset', fontweight='bold', fontsize=13)
plt.tight_layout()
plt.show()

# Print the features most correlated with the Malware target column
print("\nFeature correlations with Malware target (sorted high to low):")
print(corr_matrix['Malware'].sort_values(ascending=False).round(3))

# Choose the features most likely to differ between classes based on EDA
numeric_data['Malware'] = numeric_data['Malware'].astype(int)
top_features = [
    'SectionMaxEntropy', 'SectionMinEntropy',
    'SizeOfCode',        'NumberOfSections',
    'DllCharacteristics', 'SizeOfInitializedData'
]

fig, axes = plt.subplots(2, 3, figsize=(16, 8))
axes = axes.flatten()  # flatten to 1D for easy iteration

for i, feature in enumerate(top_features):
    sns.boxplot(
        x='Malware', y=feature, data=numeric_data,
        palette={0: '#00e676', 1: '#ff4d6d'},
        hue='Malware', legend=False,
        ax=axes[i], fliersize=2   # make outlier dots smaller for clarity
    )
    axes[i].set_title(feature, fontweight='bold')
    axes[i].set_xlabel('')
    axes[i].set_xticks([0, 1])
    axes[i].set_xticklabels(['Benign', 'Malware'])

plt.suptitle('Feature Distributions by Class: Benign vs Malware',
            fontsize=14, fontweight='bold', y=1.02)
plt.tight_layout()
plt.show()

# Confirm these columns exist before dropping (defensive check)
cols_to_drop = ['Name', 'Machine', 'TimeDateStamp', 'SectionMainChar']
print("Columns before dropping:", data.columns.tolist())

# axis=1 means "drop columns" (axis=0 would drop rows)
# inplace=True modifies the DataFrame directly without needing to reassign
data.drop(cols_to_drop, axis=1, inplace=True)

print("\nColumns after dropping:", data.columns.tolist())
print(f"\nDataset shape after dropping: {data.shape}")
print("✅ Low-value columns removed")

# Separate the feature matrix (X) from the target column (y)
# features = everything EXCEPT the target 'Malware' column
# targets  = only the 'Malware' column (0 = benign, 1 = malware)
features = data.drop("Malware", axis=1)  # shape: (n_samples, n_features)
targets  = data['Malware']               # shape: (n_samples,)

print(f"Feature matrix shape: {features.shape}")
print(f"Target vector shape:  {targets.shape}")
print(f"Feature columns: {features.columns.tolist()}")

# ── Initialise and apply MinMaxScaler ──────────────────────────────────────
# MinMaxScaler maps each feature to [0, 1] using:
#   X_scaled = (X - X_min) / (X_max - X_min)
# This ensures no single feature dominates distance calculations
scaler = MinMaxScaler()

# fit_transform() on the FULL features matrix here:
# After we split into train/test, we refit on training data only (see Section 7)
features_scaled = scaler.fit_transform(features)

print(f"\nOriginal SizeOfCode range: {features['SizeOfCode'].min():.0f} — {features['SizeOfCode'].max():.0f}")
print(f"Scaled  SizeOfCode range: 0.000000 — 1.000000")
print("✅ Features scaled to [0, 1] range")

# ── Handle class imbalance: upsample benign to minimum 10 samples ──────────
import numpy as np
from sklearn.utils import resample
from sklearn.impute import SimpleImputer

benign_mask  = targets == 0
malware_mask = targets == 1

features_benign_up, targets_benign_up = resample(
    features_scaled[benign_mask], targets[benign_mask],
    replace=True, n_samples=10, random_state=42
)
features_balanced = np.vstack([features_scaled[malware_mask], features_benign_up])
targets_balanced  = np.concatenate([targets[malware_mask], targets_benign_up])

# ── 80 / 20 train-test split ───────────────────────────────────────────────
# test_size=0.2   → 20% of data held back for testing
# stratify=targets → preserves the benign/malware ratio in both splits
# random_state=42 → fixes the random seed for reproducibility
#                   (every student running this code gets identical splits)
X_train, X_test, y_train, y_test = train_test_split(
    features_balanced, targets_balanced,
    test_size=0.2,
    stratify=targets_balanced,     # ← crucial for imbalanced datasets
    random_state=42
)

# ── Impute NaN values with median (KNN cannot handle NaNs) ────────────────
imputer = SimpleImputer(strategy='median')
X_train = imputer.fit_transform(X_train)
X_test  = imputer.transform(X_test)

print(f"Training set:  {X_train.shape[0]:>6,} files  ({X_train.shape[0]/len(targets_balanced)*100:.0f}%)")
print(f"Test set:      {X_test.shape[0]:>6,} files  ({X_test.shape[0]/len(targets_balanced)*100:.0f}%)")

# Verify that both splits maintain the same class ratio
train_malware_pct = y_train.mean() * 100
test_malware_pct  = y_test.mean()  * 100
print(f"\nMalware % in training set: {train_malware_pct:.1f}%")
print(f"Malware % in test set:     {test_malware_pct:.1f}%")
print("✅ Stratified split complete — both sets have the same class ratio")

# ── K-Nearest Neighbours ──────────────────────────────────────────────────
# n_neighbors=5: find the 5 most similar training files by Euclidean distance
# and classify by majority vote among those 5 neighbours
knn = KNeighborsClassifier(n_neighbors=5)

print("⏳ Training KNN...")
knn.fit(X_train, y_train)    # stores all training examples in memory
print("✅ KNN trained")

# Predict on the held-out test set
knn_preds = knn.predict(X_test)

knn_acc = accuracy_score(y_test, knn_preds)
print(f"\n📊 KNN Accuracy: {knn_acc:.4f} ({knn_acc*100:.2f}%)")
print("\nDetailed Classification Report:")
print(classification_report(y_test, knn_preds,
                              target_names=['Benign', 'Malware']))

# ── Random Forest ────────────────────────────────────────────────────────
# n_estimators=100: build 100 decision trees (the "forest")
# Each tree is trained on a bootstrapped sample (random rows with replacement)
# and uses a random subset of features at each split — this is the "random" part
# random_state=42: reproducible results
rfc = RandomForestClassifier(n_estimators=100, random_state=42)

print("⏳ Training Random Forest (100 trees)...")
rfc.fit(X_train, y_train)
print("✅ Random Forest trained")

rfc_preds = rfc.predict(X_test)
rfc_acc   = accuracy_score(y_test, rfc_preds)

print(f"\n📊 Random Forest Accuracy: {rfc_acc:.4f} ({rfc_acc*100:.2f}%)")
print("\nDetailed Classification Report:")
print(classification_report(y_test, rfc_preds,
                              target_names=['Benign', 'Malware']))

# ── Feature importance plot — what did the forest actually use? ───────────
feature_names = list(numeric_data.drop(columns=['Malware']).columns)
importances   = rfc.feature_importances_

plt.figure(figsize=(10, 5))
bars = plt.barh(range(len(importances)),
               sorted(importances, reverse=True),
               color='#ff6d3a')
sorted_idx = np.argsort(importances)[::-1]  # descending order of importance
plt.yticks(range(len(importances)),
           [feature_names[i] for i in sorted_idx])
plt.xlabel('Feature Importance Score')
plt.title('Random Forest Feature Importances', fontweight='bold')
plt.tight_layout()
plt.show()

# ── Logistic Regression ───────────────────────────────────────────────────
# max_iter=1000: allow up to 1000 gradient-descent steps to find the
# optimal weights. Default is 100 which may not converge on this dataset.
# random_state=42: reproducible initialisation
lr = LogisticRegression(max_iter=1000, random_state=42)

print("⏳ Training Logistic Regression...")
lr.fit(X_train, y_train)
print("✅ Logistic Regression trained")

lr_preds = lr.predict(X_test)
lr_acc   = accuracy_score(y_test, lr_preds)

print(f"\n📊 Logistic Regression Accuracy: {lr_acc:.4f} ({lr_acc*100:.2f}%)")
print("\nDetailed Classification Report:")
print(classification_report(y_test, lr_preds,
                              target_names=['Benign', 'Malware']))

# ── Decision Tree ─────────────────────────────────────────────────────────
# random_state=42: reproducible tree structure
# No max_depth set: tree will grow until all leaves are pure or contain
# fewer than min_samples_split samples — can overfit but instructive to compare
dtc = DecisionTreeClassifier(random_state=42)

print("⏳ Training Decision Tree...")
dtc.fit(X_train, y_train)
print("✅ Decision Tree trained")

dtc_preds = dtc.predict(X_test)
dtc_acc   = accuracy_score(y_test, dtc_preds)

print(f"\n📊 Decision Tree Accuracy: {dtc_acc:.4f} ({dtc_acc*100:.2f}%)")
print("\nDetailed Classification Report:")
print(classification_report(y_test, dtc_preds,
                              target_names=['Benign', 'Malware']))

# Compile results into a summary DataFrame for easy comparison
results = {
    'Algorithm':  ['KNN',     'Random Forest', 'Logistic Regression', 'Decision Tree'],
    'Model':       [knn,      rfc,              lr,                   dtc],
    'Predictions': [knn_preds, rfc_preds,        lr_preds,             dtc_preds]
}

print(f"{'Algorithm':<22} {'Accuracy':>10}")
print("─" * 35)
for name, preds in zip(results['Algorithm'], results['Predictions']):
    acc = accuracy_score(y_test, preds)
    bar = '█' * int(acc * 40)
    print(f"{name:<22} {acc*100:>6.2f}%  {bar}")

# Bar chart comparison
names  = results['Algorithm']
accs   = [accuracy_score(y_test, p) * 100 for p in results['Predictions']]
colours = ['#82aaff', '#ff6d3a', '#c3e88d', '#c792ea']

plt.figure(figsize=(9, 5))
bars = plt.bar(names, accs, color=colours, edgecolor='#1c2436', linewidth=1.2)
for bar, acc in zip(bars, accs):
    plt.text(bar.get_x() + bar.get_width()/2, bar.get_height() - 2,
             f'{acc:.2f}%', ha='center', va='top', fontweight='bold', fontsize=12)
plt.ylim(70, 101)
plt.title('Model Accuracy Comparison — Malware Detection', fontweight='bold')
plt.ylabel('Accuracy (%)')
plt.tight_layout()
plt.show()

from sklearn.metrics import ConfusionMatrixDisplay

# Plot confusion matrix for the best model (Random Forest)
fig, ax = plt.subplots(figsize=(7, 6))
cm = confusion_matrix(y_test, rfc_preds)

disp = ConfusionMatrixDisplay(confusion_matrix=cm,
                               display_labels=['Benign', 'Malware'])
disp.plot(ax=ax, cmap='Oranges', colorbar=False)

ax.set_title('Random Forest — Confusion Matrix', fontweight='bold', fontsize=13)
print("\nInterpretation:")
print(f"  True Negatives  (Benign caught correctly):   {cm[0,0]:,}")
print(f"  False Positives (Benign wrongly flagged):     {cm[0,1]:,}")
print(f"  False Negatives (Malware that slipped through):{cm[1,0]:,}")
print(f"  True Positives  (Malware caught correctly):   {cm[1,1]:,}")
plt.tight_layout()
plt.show()

from sklearn.metrics import roc_curve, auc

# predict_proba() returns [P(benign), P(malware)] for each file
# We take column [:, 1] — the malware probability — as the score for ROC
models_for_roc = [
    ('KNN',                 knn),
    ('Random Forest',       rfc),
    ('Logistic Regression', lr),
    ('Decision Tree',       dtc),
]
colours = ['#82aaff', '#ff6d3a', '#c3e88d', '#c792ea']

plt.figure(figsize=(8, 7))

for (name, model), colour in zip(models_for_roc, colours):
    proba        = model.predict_proba(X_test)[:, 1]
    fpr, tpr, _  = roc_curve(y_test, proba)
    roc_auc      = auc(fpr, tpr)
    plt.plot(fpr, tpr, color=colour, lw=2,
             label=f'{name} (AUC = {roc_auc:.4f})')

# Diagonal = random classifier baseline (AUC = 0.50)
plt.plot([0, 1], [0, 1], '--', color='#4a5470',
         label='Random baseline (AUC = 0.5000)')

plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.02])
plt.xlabel('False Positive Rate  (benign files wrongly flagged)')
plt.ylabel('True Positive Rate   (malware files caught)')
plt.title('ROC Curves — All Models', fontweight='bold')
plt.legend(loc='lower right')
plt.tight_layout()
plt.show()

import pickle

# ── Save the trained Random Forest model ──────────────────────────────────
# 'wb' = write binary — pickle produces binary data, not text
model_file = "malware_predictor.pkl"

with open(model_file, "wb") as file:
    pickle.dump(rfc, file)   # serialise the RandomForestClassifier object

print(f"✅ Model saved to '{model_file}'")

# ── Save the MinMaxScaler (must be saved alongside the model) ─────────────
# The scaler stores the min/max of each feature from the training set.
# New data must be scaled with these exact same values — never refit the scaler.
scaler_file = "malware_scaler.pkl"

with open(scaler_file, "wb") as file:
    pickle.dump(scaler, file)

print(f"✅ Scaler saved to '{scaler_file}'")
print("\n📌 Both files will appear in the Colab Files panel (left sidebar).")
print("   Download them to your machine to use this model outside of Colab.")

# ── Load the saved model from disk ────────────────────────────────────────
# 'rb' = read binary
with open("malware_predictor.pkl", "rb") as file:
    rfc_loaded = pickle.load(file)

with open("malware_scaler.pkl", "rb") as file:
    scaler_loaded = pickle.load(file)

print(f"✅ Loaded model type: {type(rfc_loaded)}")
print(f"   Number of trees in forest: {len(rfc_loaded.estimators_)}")

# Verify the loaded model gives identical predictions to the original
loaded_preds = rfc_loaded.predict(X_test)
match = (loaded_preds == rfc_preds).all()
print(f"\n{'✅' if match else '❌'} Loaded model predictions {'match' if match else 'do NOT match'} original")
print(f"   Accuracy of loaded model: {accuracy_score(y_test, loaded_preds)*100:.2f}%")

# ── Example: classify a new, hypothetical PE file ─────────────────────────
# In a real deployment, these values would be extracted from an actual
# .exe or .dll file using a PE parsing library like 'pefile' (pip install pefile)
# Here we construct a sample by hand to demonstrate the prediction pipeline

# Feature order must exactly match the order used during training
# (same as the columns in the DataFrame after dropping 4 columns)
feature_order = list(numeric_data.drop(columns=['Malware']).columns)  # e.g. ['SizeOfOptionalHeader', 'SizeOfCode', ...]

# Example 1: suspicious file (high entropy, few sections, small code size)
suspicious_file = {
    col: 0.0 for col in feature_order  # start all at 0
}
# Manually set suspicious values for known malware indicators
suspicious_file['SectionMaxEntropy']  = 7.8   # very high — packed/encrypted
suspicious_file['SectionMinEntropy']  = 7.5   # very high
suspicious_file['NumberOfSections']  = 2     # unusually few sections
suspicious_file['SizeOfCode']        = 4096  # tiny code section = packed

def predict_file(file_features, scaler, model):
    """Scale features and run prediction. Returns 'MALWARE' or 'BENIGN'."""
    raw = np.array([[file_features[col] for col in feature_order]])
    scaled = scaler.transform(raw)    # ← .transform(), NOT fit_transform()
    pred   = model.predict(scaled)[0]
    proba  = model.predict_proba(scaled)[0]
    label  = '🔴 MALWARE' if pred == 1 else '🟢 BENIGN'
    print(f"  Prediction:  {label}")
    print(f"  Probability: Benign={proba[0]:.3f}  Malware={proba[1]:.3f}")
    return pred

print("── Suspicious file (high entropy, tiny code) ──")
predict_file(suspicious_file, scaler_loaded, rfc_loaded)

# ══════════════════════════════════════════════════════════════════════════
# MALWARE SCANNER — Complete Colab launcher
# ══════════════════════════════════════════════════════════════════════════

# ── Step 1: Install dependencies ──────────────────────────────────────────
!pip install streamlit pefile plotly pyngrok -q

# ── Step 2: Verify model files exist ──────────────────────────────────────
import os
for f in ["malware_predictor.pkl", "malware_scaler.pkl"]:
    if os.path.exists(f):
        print(f"✅ Found: {f}")
    else:
        raise FileNotFoundError(f"❌ Missing: {f} — upload it first")

# ── Step 3: Write scanner_app.py ──────────────────────────────────────────
app_code = '''
import streamlit as st
import pickle
import numpy as np
import pandas as pd
import pefile
import plotly.graph_objects as go
import math

st.set_page_config(page_title="Malware Scanner", page_icon="🛡️", layout="centered")

@st.cache_resource
def load_model():
    with open("malware_predictor.pkl", "rb") as f:
        model = pickle.load(f)
    with open("malware_scaler.pkl", "rb") as f:
        scaler = pickle.load(f)
    return model, scaler

model, scaler = load_model()

FEATURE_NAMES = [
    "e_magic", "e_cblp", "e_cp", "e_crlc", "e_cparhdr",
    "e_minalloc", "e_maxalloc", "e_ss", "e_sp", "e_csum",
    "e_ip", "e_cs", "e_lfarlc", "e_ovno", "e_oemid",
    "e_oeminfo", "e_lfanew", "NumberOfSections",
    "PointerToSymbolTable", "NumberOfSymbols",
    "SizeOfOptionalHeader", "Characteristics", "Magic",
    "MajorLinkerVersion", "MinorLinkerVersion", "SizeOfCode",
    "SizeOfInitializedData", "SizeOfUninitializedData",
    "AddressOfEntryPoint", "BaseOfCode", "ImageBase",
    "SectionAlignment", "FileAlignment",
    "MajorOperatingSystemVersion", "MinorOperatingSystemVersion",
    "MajorImageVersion", "MinorImageVersion",
    "MajorSubsystemVersion", "MinorSubsystemVersion",
    "SizeOfHeaders", "CheckSum", "SizeOfImage", "Subsystem",
    "DllCharacteristics", "SizeOfStackReserve", "SizeOfStackCommit",
    "SizeOfHeapReserve", "SizeOfHeapCommit", "LoaderFlags",
    "NumberOfRvaAndSizes", "SuspiciousImportFunctions",
    "SuspiciousNameSection", "SectionsLength",
    "SectionMinEntropy", "SectionMaxEntropy",
    "SectionMinRawsize", "SectionMaxRawsize",
    "SectionMinVirtualsize", "SectionMaxVirtualsize",
    "SectionMaxPhysical", "SectionMinPhysical",
    "SectionMaxVirtual", "SectionMinVirtual",
    "SectionMaxPointerData", "SectionMinPointerData",
    "SectionMaxChar", "DirectoryEntryImport",
    "DirectoryEntryImportSize", "DirectoryEntryExport",
    "ImageDirectoryEntryExport", "ImageDirectoryEntryImport",
    "ImageDirectoryEntryResource", "ImageDirectoryEntryException",
    "ImageDirectoryEntrySecurity",
]

SUSPICIOUS_IMPORTS = [
    "VirtualAlloc", "VirtualProtect", "WriteProcessMemory",
    "CreateRemoteThread", "ShellExecute", "WinExec",
    "RegSetValue", "URLDownloadToFile", "socket", "connect",
    "CreateProcess", "OpenProcess", "NtUnmapViewOfSection"
]

def get_entropy(data):
    if not data:
        return 0.0
    freq = {}
    for b in data:
        freq[b] = freq.get(b, 0) + 1
    entropy = 0.0
    for count in freq.values():
        p = count / len(data)
        entropy -= p * math.log2(p)
    return entropy

def extract_features(file_bytes):
    features = {name: 0.0 for name in FEATURE_NAMES}
    try:
        pe = pefile.PE(data=file_bytes)
        dos = pe.DOS_HEADER
        for field in ["e_magic","e_cblp","e_cp","e_crlc","e_cparhdr",
                      "e_minalloc","e_maxalloc","e_ss","e_sp","e_csum",
                      "e_ip","e_cs","e_lfarlc","e_ovno","e_oemid",
                      "e_oeminfo","e_lfanew"]:
            features[field] = getattr(dos, field, 0)
        fh = pe.FILE_HEADER
        features["NumberOfSections"]     = fh.NumberOfSections
        features["PointerToSymbolTable"] = fh.PointerToSymbolTable
        features["NumberOfSymbols"]      = fh.NumberOfSymbols
        features["SizeOfOptionalHeader"] = fh.SizeOfOptionalHeader
        features["Characteristics"]      = fh.Characteristics
        if hasattr(pe, "OPTIONAL_HEADER"):
            opt = pe.OPTIONAL_HEADER
            for field in ["Magic","MajorLinkerVersion","MinorLinkerVersion",
                          "SizeOfCode","SizeOfInitializedData","SizeOfUninitializedData",
                          "AddressOfEntryPoint","BaseOfCode","ImageBase",
                          "SectionAlignment","FileAlignment",
                          "MajorOperatingSystemVersion","MinorOperatingSystemVersion",
                          "MajorImageVersion","MinorImageVersion",
                          "MajorSubsystemVersion","MinorSubsystemVersion",
                          "SizeOfHeaders","CheckSum","SizeOfImage","Subsystem",
                          "DllCharacteristics","SizeOfStackReserve","SizeOfStackCommit",
                          "SizeOfHeapReserve","SizeOfHeapCommit","LoaderFlags",
                          "NumberOfRvaAndSizes"]:
                features[field] = getattr(opt, field, 0)
            if hasattr(opt, "DATA_DIRECTORY"):
                dd = opt.DATA_DIRECTORY
                features["ImageDirectoryEntryExport"]    = dd[0].Size if len(dd) > 0 else 0
                features["ImageDirectoryEntryImport"]    = dd[1].Size if len(dd) > 1 else 0
                features["ImageDirectoryEntryResource"]  = dd[2].Size if len(dd) > 2 else 0
                features["ImageDirectoryEntryException"] = dd[3].Size if len(dd) > 3 else 0
                features["ImageDirectoryEntrySecurity"]  = dd[4].Size if len(dd) > 4 else 0
        if pe.sections:
            entropies  = [get_entropy(s.get_data()) for s in pe.sections]
            raw_sizes  = [s.SizeOfRawData           for s in pe.sections]
            virt_sizes = [s.Misc_VirtualSize         for s in pe.sections]
            phys       = [s.PointerToRawData         for s in pe.sections]
            virt       = [s.VirtualAddress           for s in pe.sections]
            ptr_data   = [s.PointerToRelocations     for s in pe.sections]
            chars      = [s.Characteristics          for s in pe.sections]
            features["SectionMinEntropy"]     = min(entropies)
            features["SectionMaxEntropy"]     = max(entropies)
            features["SectionMinRawsize"]     = min(raw_sizes)
            features["SectionMaxRawsize"]     = max(raw_sizes)
            features["SectionMinVirtualsize"] = min(virt_sizes)
            features["SectionMaxVirtualsize"] = max(virt_sizes)
            features["SectionMaxPhysical"]    = max(phys)
            features["SectionMinPhysical"]    = min(phys)
            features["SectionMaxVirtual"]     = max(virt)
            features["SectionMinVirtual"]     = min(virt)
            features["SectionMaxPointerData"] = max(ptr_data)
            features["SectionMinPointerData"] = min(ptr_data)
            features["SectionMaxChar"]        = max(chars)
            features["SectionsLength"]        = len(pe.sections)
            suspicious_names = [".upx", ".aspack", ".adata", "UPX"]
            features["SuspiciousNameSection"] = sum(
                1 for s in pe.sections
                if s.Name.decode(errors="ignore").strip("\\x00") in suspicious_names
            )
        if hasattr(pe, "DIRECTORY_ENTRY_IMPORT"):
            imports = pe.DIRECTORY_ENTRY_IMPORT
            features["DirectoryEntryImport"]     = 1
            features["DirectoryEntryImportSize"] = len(imports)
            all_imports = []
            for entry in imports:
                for imp in entry.imports:
                    if imp.name:
                        all_imports.append(imp.name.decode(errors="ignore"))
            features["SuspiciousImportFunctions"] = sum(
                1 for fn in all_imports if fn in SUSPICIOUS_IMPORTS
            )
        if hasattr(pe, "DIRECTORY_ENTRY_EXPORT"):
            features["DirectoryEntryExport"] = 1
        pe.close()
    except Exception as e:
        st.warning(f"⚠️ Feature extraction partial: {e}")
    return features

# ── UI ─────────────────────────────────────────────────────────────────────
st.title("🛡️ Malware Scanner")
st.caption("Upload a Windows PE executable (.exe or .dll) for instant static analysis. The file is never executed.")

uploaded_file = st.file_uploader("Drop your file here", type=["exe", "dll"])

if uploaded_file:
    file_bytes = uploaded_file.read()
    st.info(f"📁 **{uploaded_file.name}** — {len(file_bytes)/1024:.1f} KB")
    with st.spinner("Extracting PE features..."):
        features = extract_features(file_bytes)
    feature_vector = np.array([[features[f] for f in FEATURE_NAMES]])
    scaled         = scaler.transform(feature_vector)
    prediction     = model.predict(scaled)[0]
    proba          = model.predict_proba(scaled)[0]
    st.divider()
    if prediction == 1:
        st.error("## 🔴 MALWARE DETECTED")
    else:
        st.success("## 🟢 BENIGN")
    col1, col2 = st.columns(2)
    col1.metric("Benign probability",  f"{proba[0]*100:.1f}%")
    col2.metric("Malware probability", f"{proba[1]*100:.1f}%")
    st.divider()
    st.subheader("📊 Feature Breakdown")
    importances = model.feature_importances_
    top_idx     = np.argsort(importances)[::-1][:15]
    fig = go.Figure(go.Bar(
        x=[features[FEATURE_NAMES[i]] for i in top_idx],
        y=[FEATURE_NAMES[i]           for i in top_idx],
        orientation="h",
        marker_color=["#ff4d6d" if prediction == 1 else "#00e676"] * 15
    ))
    fig.update_layout(
        title="Top 15 Features (by model importance)",
        xaxis_title="Raw Feature Value",
        yaxis=dict(autorange="reversed"),
        height=500,
        margin=dict(l=10, r=10, t=40, b=10)
    )
    st.plotly_chart(fig, use_container_width=True)
    with st.expander("🔍 All extracted features"):
        df = pd.DataFrame(features.items(), columns=["Feature", "Value"])
        st.dataframe(df, use_container_width=True)
'''

with open("scanner_app.py", "w") as f:
    f.write(app_code)
print("✅ scanner_app.py written")

# ── Step 4: Launch Streamlit + ngrok tunnel ───────────────────────────────
import subprocess, time
from pyngrok import ngrok

NGROK_TOKEN = "your_authtoken_here"   # ← paste your token from dashboard.ngrok.com
ngrok.set_auth_token(NGROK_TOKEN)
ngrok.kill()

process = subprocess.Popen(
    ["streamlit", "run", "scanner_app.py",
     "--server.port", "8501",
     "--server.headless", "true",
     "--server.enableCORS", "false"],
    stdout=subprocess.DEVNULL,
    stderr=subprocess.DEVNULL
)
time.sleep(4)

public_url = ngrok.connect(8501)
print(f"\n{'='*55}")
print(f"  🛡️  MALWARE SCANNER IS LIVE")
print(f"  👉  {public_url}")
print(f"{'='*55}")
print("\n  Test with: C:\\Windows\\System32\\notepad.exe")