Alina-sul/malware-detection-ml

# 恶意软件检测软件 使用 CICMalMem2022 内存分析数据集的恶意软件检测系统，包含 MLflow 实验跟踪和 Streamlit UI。 ## 数据集 **来源：** [CICMalMem2022](https://www.unb.ca/cic/datasets/malmem-2022.html) (来自 Kaggle 的混淆变体) 该数据集包含使用 VolMemLyzer（Volatility 封装器）提取的 **58,058 个** 内存转储特征样本： - **29,227 个** 良性样本 - **28,831 个** 跨 3 个家族的恶意软件样本： - **特洛伊木马** (5 种类型): Zeus, Emotet, Refroso, Scar, Reconyc - **间谍软件** (5 种类型): 180Solutions, CWS, Gator, Transponder, TIBS - **勒索软件** (5 种类型): Conti, MAZE, Pysa, Ako, Shade 来自 Volatility 插件的 **55 个数值特征**：pslist, dlllist, handles, ldrmodules, malfind, psxview, modules, svcscan, callbacks。 **标签列：** - `Class` — 二分类标签：`Benign` / `Malware` - `Category` — 源文件名（包含恶意软件家族信息，例如 `Ransomware-Ako-...`） 无空值。数据为 parquet 格式。 ## 项目结构 ``` ├── config/ │ ├── config.yaml # All settings: paths, hyperparams, MLflow │ └── selected_features.yaml # Manually curated feature subset (tracked) ├── data/ │ ├── raw/ # Downloaded dataset (gitignored) │ ├── processed/ # Train/val/test splits (gitignored) │ ├── README.md # Dataset stats, features, split info │ └── GLOSSARY.md # Malware variants & OS concepts ├── docs/ │ ├── AI_USAGE.md # Academic AI usage disclosure │ └── FEATURE_SELECTION.md # Feature selection methods guide & glossary ├── models/ # Trained model artifacts (tracked) │ ├── selected/ # Models trained on the 24-feature subset │ └── all/ # Models trained on all 55 features ├── src/ │ ├── cli/ │ │ ├── commands.py # CLI group root │ │ ├── data_commands.py # split-data, compare-features, dashboard │ │ └── train_commands.py # train command (thin wrapper over runner) │ ├── data/ │ │ ├── loader.py # Dataset loading & family extraction │ │ ├── preprocessing.py # Label encoding, train/val/test split │ │ ├── feature_info.py # Parser for per-feature descriptions │ │ └── feature_engineering/ # Feature selection framework │ │ ├── helpers.py # Shared utilities (ranking builder, feature groups) │ │ ├── scaling.py # Standard / MinMax scaling │ │ ├── ranking.py # 5 ranking methods (variance, correlation, ANOVA, MI, RF) │ │ ├── consensus.py # Borda-count consensus ranking │ │ ├── reporting.py # Plots & report generation │ │ └── selector.py # FeatureSelector class (orchestrator, incl. "manual" mode) │ ├── models/ # Model wrappers — one file per classifier │ │ ├── base.py # BaseModel ABC + SklearnModel + ScaledSklearnModel │ │ ├── dummy.py # DummyClassifier (majority-class baseline) │ │ ├── logistic_regression.py # LR with StandardScaler (linear baseline) │ │ ├── random_forest.py # RandomForestClassifier (tree bagging) │ │ ├── xgboost_model.py # XGBoost via Booster API (stable on macOS ARM) │ │ ├── dnn.py # PyTorch feed-forward net with early stopping │ │ └── dnn_helpers.py # Pure DNN helpers (build_network, train_one_epoch, ...) │ ├── training/ # Orchestrator, metrics, MLflow logging │ │ ├── runner.py # run_training() — shared by CLI and dashboard │ │ ├── trainer.py # train_model / train_all │ │ ├── evaluate.py # compute_metrics (acc, precision, recall, F1, AUC) │ │ └── experiment.py # MLflow logger (log_run, configure_mlflow) │ ├── inference/ # Prediction, VolMemLyzer integration (Phase 3) │ └── utils/ # Config loader, plotting, logging ├── app/ # Streamlit web UI │ ├── streamlit_app.py # Entry point & welcome page │ ├── helpers/ │ │ ├── data.py # Cached loaders (config, dataset, comparison) │ │ ├── charts.py # Plotly chart builders (including TP/FP/TN/FN labelled confusion matrix) │ │ ├── layout.py # Reusable layout components │ │ ├── selection.py # Save/load manually-curated selection YAML │ │ ├── feature_builder.py # Pure helpers for the Feature Builder page │ │ └── training.py # Pure helpers for the Model Dashboard page │ └── pages/ │ ├── 1_Dataset_Overview.py │ ├── 2_Feature_Selection.py # Interactive feature-selection comparison │ ├── 3_Feature_Builder.py # Curate the final feature subset │ ├── 4_Model_Dashboard.py # Compare models + in-page training launcher │ └── 5_Predict.py # Coming in Phase 3 (inference) ├── reports/ │ ├── feature_selection/ # Static feature comparison report (tracked) │ └── training/ # JSON training summaries (tracked) ├── main.py # CLI entry point └── tests/ # pytest suite (199+ tests) ``` ## 设置 ``` python3.12 -m venv .venv source .venv/bin/activate pip install -r requirements.txt ``` ## 使用方法 所有管道操作均通过 CLI 运行： ``` python main.py --help # Show available commands python main.py split-data # Split dataset into train/val/test (60/20/20) python main.py compare-features # Run feature selection comparison & generate static report python main.py train # Train all 5 classifiers on the 24-feature subset python main.py train --features all # ...or on all 55 features python main.py train --model xgboost # ...or just one model python main.py dashboard # Launch the interactive Streamlit dashboard ``` ### 特征选择 本项目中的特征选择分三个阶段进行： **1. 对比（统计，建模前）。** `compare-features` CLI 命令（或仪表盘中的 **Feature Selection** 页面）在模型训练前运行五种方法对所有 55 个特征进行排名： - **Variance Threshold** — 标记近似常量的特征 - **Correlation Analysis** — 识别冗余特征对 - **ANOVA F-test** — 衡量线性类别可分性 - **Mutual Information** — 捕获非线性依赖关系 - **RF Importance** — 来自轻量级随机森林的 Gini 重要性 静态结果保存至 `reports/feature_selection/`（共识排名 CSV、各方法 JSON、相关性热图、排名图、文本摘要）。 **2. 人工整理（提交最终子集）。** 仪表盘中的 **Feature Builder** 页面在一个可编辑的表格中展示所有 55 个特征，并标注了共识排名、方法间分歧、插件组和警告标志。默认为按共识排名的前 25 个特征；可以交互式地添加或移除特征。点击 **Save selection** 会将列表持久化到 `config/selected_features.yaml` —— 这是一个受跟踪的产物，作为最终特征子集的唯一真实来源。 **3. 训练时的应用。** 在 `config/config.yaml` 中设置 `features.selection_method: "manual"`，`FeatureSelector.apply()` 管道会读取 `config/selected_features.yaml`，并在模型看到数据之前将训练 / 验证 / 测试划分过滤到提交的列。 请参阅 [docs/FEATURE_SELECTION.md](docs/FEATURE_SELECTION.md) 以获取关于解读分数、阈值和选择最终特征子集的指南。 ### 模型训练 比较了五个分类器，形成了一个从随机猜测到深度学习的可辩护的复杂度阶梯： 1. **Dummy**（多数类基线）—— 绝对下限；在这个平衡数据集上，任何真实模型都必须击败约 50% 的准确率。 2. **Logistic Regression** —— 线性科学对照；如果树 / 神经模型没有明显击败 LR，则该任务本质上是线性的。 3. **Random Forest** —— 树 bagging（项目规范要求）。 4. **XGBoost** —— 树 boosting（项目规范要求）。 5. **Deep Neural Network** —— PyTorch 前馈网络，带有早停机制（项目规范要求）。 训练由 `python main.py train` 或 Model Dashboard 页面内的 **Train models** 启动器驱动（两者都调用相同的 `src/training/runner.py::run_training` 函数）。每次运行： - 应用配置的特征选择（`selected` → 24 个提交的特征，`all` → 每个特征） - 使用 `config/config.yaml` 中的固定超参数拟合每个请求的模型 - 使用 `accuracy`、`precision`、`recall`、`F1`、`ROC-AUC` 和完整的 2×2 混淆矩阵在验证和测试划分上进行评估（scikit-learn 默认值，请参阅 [scoring API reference](https://scikit-learn.org/stable/modules/model_evaluation.html#scoring-api-overview)） - 将每个训练好的模型保存到 `models/{feature_set}/{name}.{ext}`（`.joblib` / `.pt` / `.json` 取决于底层库） - 使用超参数、指标、混淆矩阵和保存的模型文件将每次运行记录到 MLflow 的 `malware-detection` 实验下 - 将紧凑的 JSON 结果摘要写入 `reports/training/training_summary_{feature_set}.json` 在 Streamlit 应用的 **Model Dashboard** 页面中并排浏览训练好的模型：按指标对比 selected 与 all 特征集的分组条形图、合并的结果表、带有明确 TN/FP/FN/TP 标签的每个模型混淆矩阵，以及一个页内训练启动器，允许您在不离开浏览器的情况下使用不同的模型 / 特征集组合重新训练。 原始 MLflow 运行数据位于 `mlruns/` 下（已 gitignore）。启动 `mlflow ui --backend-store-uri mlruns/` 以直接探索它们。 ## 技术栈 - **ML：** scikit-learn, XGBoost, PyTorch - **实验跟踪：** MLflow - **Web UI：** Streamlit - **内存分析：** Volatility 3, VolMemLyzer

标签：Apex, BSD, CICMalMem2022, JARM, Kubernetes, MLflow, Parquet, Streamlit, 内存分析, 凭据扫描, 分类任务, 勒索软件, 实验跟踪, 数据挖掘, 木马检测, 机器学习, 模型部署, 特征工程, 特征选择, 网络安全, 访问控制, 逆向工具, 间谍软件, 隐私保护