labs/Lab-5/Lab-5-Solution/06_Capstone_Malware_Detection_Solution.ipynb

Module 5 — Lab 6: Capstone — An Explainable Malware Detection Workflow

Topics: end-to-end workflow, model comparison, error evidence, analyst recommendation

Solution version: Run the cells in order, compare the outputs with your work, and read the suggested answers.

Mission

Build and report a malware detector using the prepared UCI feature table. Compare an explainable bounded tree with a deeper tree and a Random Forest, then recommend a model based on detection quality, false positives, interpretability, and dataset limitations.

Required outputs

• output/capstone_metrics.csv
• output/capstone_errors.csv
• output/capstone_tree_rules.txt
• output/capstone_confusion_matrix.png
• output/capstone_analyst_summary.md

In [ ]:

from pathlib import Path
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

DATA_PATH = Path("data/malware_features.csv")

if not DATA_PATH.exists():
    raise FileNotFoundError(
        "Missing data/malware_features.csv. Follow data/README.md, place the three "
        "official UCI CSV files under data/raw/, and run: python data/prepare_dataset.py"
    )


from sklearn.dummy import DummyClassifier
from sklearn.impute import SimpleImputer
from sklearn.metrics import confusion_matrix, precision_score, recall_score, f1_score, accuracy_score
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.tree import DecisionTreeClassifier


from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import ConfusionMatrixDisplay
from sklearn.tree import export_text


def security_metrics(y_true, y_pred):
    tn, fp, fn, tp = confusion_matrix(y_true, y_pred, labels=[0, 1]).ravel()
    return {
        "accuracy": accuracy_score(y_true, y_pred),
        "malware_precision": precision_score(y_true, y_pred, zero_division=0),
        "malware_recall": recall_score(y_true, y_pred, zero_division=0),
        "malware_f1": f1_score(y_true, y_pred, zero_division=0),
        "false_positives": int(fp),
        "false_negatives": int(fn),
        "false_positive_rate": fp / (fp + tn) if fp + tn else 0.0,
    }


OUTPUT_DIR = Path("output")
OUTPUT_DIR.mkdir(exist_ok=True)

In [ ]:

df = pd.read_csv(DATA_PATH, low_memory=False)
feature_cols = [column for column in df.columns if column not in {"source", "label"}]
X = df[feature_cols].apply(pd.to_numeric, errors="coerce")
y = df["label"].astype(int)
X_train, X_test, y_train, y_test, source_train, source_test = train_test_split(
    X, y, df["source"], test_size=0.20, stratify=y, random_state=42
)
print("Train/test rows:", len(X_train), len(X_test))

1. Build and compare candidate models

In [ ]:

models = {
    "Dummy": DummyClassifier(strategy="most_frequent"),
    "ExplainableTree": Pipeline([
        ("imputer", SimpleImputer(strategy="median", add_indicator=True)),
        ("tree", DecisionTreeClassifier(max_depth=4, min_samples_leaf=10, random_state=42)),
    ]),
    "DeepTree": Pipeline([
        ("imputer", SimpleImputer(strategy="median", add_indicator=True)),
        ("tree", DecisionTreeClassifier(random_state=42)),
    ]),
    "RandomForest": Pipeline([
        ("imputer", SimpleImputer(strategy="median", add_indicator=True)),
        ("forest", RandomForestClassifier(
            n_estimators=200,
            min_samples_leaf=2,
            random_state=42,
            n_jobs=-1,
        )),
    ]),
}

metric_rows = []
predictions = {}
probabilities = {}
for name, model in models.items():
    model.fit(X_train, y_train)
    pred = model.predict(X_test)
    predictions[name] = pred
    if hasattr(model, "predict_proba"):
        probabilities[name] = model.predict_proba(X_test)[:, 1]
    metric_rows.append({"model": name, **security_metrics(y_test, pred)})

metrics_table = pd.DataFrame(metric_rows).set_index("model").sort_values(
    ["malware_recall", "malware_f1"], ascending=False
)
display(metrics_table.round(4))

2. Select a model with an explicit decision rule

Prefer the explainable tree when its malware recall is within 0.03 of the best non-dummy recall and its false-positive rate is no more than 0.05 above the best non-dummy FPR. Otherwise choose the non-dummy model with the highest malware F1.

In [ ]:

eligible_metrics = metrics_table.drop(index="Dummy")
best_recall = eligible_metrics["malware_recall"].max()
best_fpr = eligible_metrics["false_positive_rate"].min()
explainable_ok = (
    eligible_metrics.loc["ExplainableTree", "malware_recall"] >= best_recall - 0.03
    and eligible_metrics.loc["ExplainableTree", "false_positive_rate"] <= best_fpr + 0.05
)
selected_name = (
    "ExplainableTree"
    if explainable_ok
    else eligible_metrics["malware_f1"].idxmax()
)
selected_model = models[selected_name]
selected_predictions = predictions[selected_name]
print("Selected model:", selected_name)

3. Create error evidence and a confusion matrix

In [ ]:

selected_probability = probabilities.get(selected_name, selected_predictions.astype(float))
error_table = pd.DataFrame({
    "source": source_test.to_numpy(),
    "true_label": y_test.to_numpy(),
    "predicted_label": selected_predictions,
    "malware_probability": selected_probability,
}, index=X_test.index)
error_table["error_type"] = np.select(
    [
        (error_table["true_label"] == 1) & (error_table["predicted_label"] == 0),
        (error_table["true_label"] == 0) & (error_table["predicted_label"] == 1),
    ],
    ["false_negative", "false_positive"],
    default="correct",
)
errors_only = error_table[error_table["error_type"] != "correct"].copy()
display(errors_only["error_type"].value_counts())

fig, ax = plt.subplots(figsize=(6, 5))
ConfusionMatrixDisplay.from_predictions(
    y_test,
    selected_predictions,
    labels=[0, 1],
    display_labels=["benign", "malware"],
    cmap="Blues",
    ax=ax,
    colorbar=False,
)
ax.set_title(f"{selected_name} — held-out test set")
fig.tight_layout()
fig.savefig(OUTPUT_DIR / "capstone_confusion_matrix.png", dpi=160)
plt.show()

4. Export rules and write the analyst recommendation

In [ ]:

explainable_model = models["ExplainableTree"]
explainable_imputer = explainable_model.named_steps["imputer"]
explainable_estimator = explainable_model.named_steps["tree"]
explainable_names = explainable_imputer.get_feature_names_out(feature_cols)
tree_rules = export_text(
    explainable_estimator,
    feature_names=list(explainable_names),
    max_depth=4,
)

selected_metrics = metrics_table.loc[selected_name]
analyst_summary = "\n".join([
    "# Malware Detection Capstone — Analyst Summary",
    "",
    "## Recommendation",
    "",
    f"Selected model: **{selected_name}**.",
    "",
    f"- Malware recall: {selected_metrics['malware_recall']:.3f}",
    f"- Malware precision: {selected_metrics['malware_precision']:.3f}",
    f"- Malware F1: {selected_metrics['malware_f1']:.3f}",
    f"- False positives: {int(selected_metrics['false_positives'])}",
    f"- False negatives: {int(selected_metrics['false_negatives'])}",
    "",
    "The selection followed the documented recall/FPR/interpretability rule. "
    "The bounded tree is preferred when its detection quality is close to the "
    "strongest candidate; otherwise the best non-dummy F1 model is selected.",
    "",
    "## Operational interpretation",
    "",
    "False negatives represent malware that would pass the detector. False positives "
    "represent benign files that would create review workload. A deployment threshold "
    "must therefore be matched to analyst capacity and the cost of a missed infection.",
    "",
    "## Limitations",
    "",
    "This historical dataset is imbalanced, and benign/malware records come from "
    "different source collections. Random-split performance may include collection "
    "artifacts and does not demonstrate robustness to future malware, concept drift, "
    "packing, or adversarial evasion. The model should be used as one triage signal "
    "alongside signatures, sandboxing, and analyst review.",
    "",
])

metrics_table.to_csv(OUTPUT_DIR / "capstone_metrics.csv")
errors_only.to_csv(OUTPUT_DIR / "capstone_errors.csv", index=True)
(OUTPUT_DIR / "capstone_tree_rules.txt").write_text(tree_rules, encoding="utf-8")
(OUTPUT_DIR / "capstone_analyst_summary.md").write_text(analyst_summary, encoding="utf-8")
print(analyst_summary)

Self-check

In [ ]:

expected_files = {
    "capstone_metrics.csv",
    "capstone_errors.csv",
    "capstone_tree_rules.txt",
    "capstone_confusion_matrix.png",
    "capstone_analyst_summary.md",
}
assert set(models) == {"Dummy", "ExplainableTree", "DeepTree", "RandomForest"}
assert selected_name in set(models) - {"Dummy"}
assert len(error_table) == len(y_test)
assert set(errors_only["error_type"].unique()).issubset({"false_negative", "false_positive"})
assert isinstance(tree_rules, str) and "class:" in tree_rules
assert expected_files.issubset({path.name for path in OUTPUT_DIR.iterdir()})
print("✓ Lab 6 completed")

Suggested review checklist

• The recommendation follows a rule stated before inspecting individual errors.
• Accuracy is not used as the sole decision metric.
• False negatives and false positives are translated into operational consequences.
• source is used only for auditing, never as a model feature.
• Tree rules are treated as model explanations, not universal malware signatures.
• Dataset age, imbalance, source confounding, drift, and evasion are acknowledged.
• The detector is positioned as one layer in a defense workflow, not an autonomous verdict.