Skip to content

Surajit-Samanta/MalwareLens_AI

Repository files navigation

MalwareLens AI

A static-analysis malware triage workbench built for AICS-108 (Intelligent Malware Analysis). It takes engineered file features and runs them through a Random Forest baseline and a scikit-learn MLP to flag samples as malicious-like or benign-like, assign a P1/P2/P3 priority, map high-risk signals to MBC/MITRE ATT&CK behavior categories, cluster samples into likely campaigns, and auto-draft a reviewable YARA-style detection rule.

Safety note: this project never downloads, stores, or executes live malware. The standard pipeline runs entirely on synthetic feature metadata. The advanced extension uses pre-extracted static feature vectors from the public EMBER dataset — never raw executables.

Features

  • Synthetic corpus generation — 2,400 synthetic records across 6 families (benign + 5 malicious-like families), 22 engineered features per record.
  • Dual-model training — Random Forest baseline + scikit-learn deep MLP, both wrapped in StandardScaler pipelines, with metrics/confusion-matrix output.
  • Single-sample triage — classification, malicious probability, priority band, and MBC/ATT&CK-mapped behavioral signals for one sample at a time.
  • Campaign clustering — PCA + KMeans over the feature space to group samples into likely campaigns, with a saved scatter plot and cluster/family breakdown.
  • AI-assisted YARA rule generation — drafts a reviewable .yar rule from a sample's flavor strings (explicitly marked as requiring human review before operational use).
  • Real-data cross-check (advanced extension) — maps EMBER2018/2024 raw feature records into this project's schema so the same training/eval pipeline can be run against real static features, not just synthetic data.

Project layout

MalwareLens_AI/
├── data
│   ├── malwarelens_synthetic_features.csv
│   ├── malwarelens_synthetic_features.csv:Zone.Identifier
│   ├── malwarelens_synthetic_features.csv:Zone.Identifier:Zone.Identifier
│   └── safe_samples
│       ├── benign_admin.json
│       ├── benign_admin.json:Zone.Identifier
│       ├── benign_admin.json:Zone.Identifier:Zone.Identifier
│       ├── botnet_like.json
│       ├── botnet_like.json:Zone.Identifier
│       ├── botnet_like.json:Zone.Identifier:Zone.Identifier
│       ├── dropper_like.json
│       ├── dropper_like.json:Zone.Identifier
│       ├── dropper_like.json:Zone.Identifier:Zone.Identifier
│       ├── infostealer_like.json
│       ├── infostealer_like.json:Zone.Identifier
│       ├── infostealer_like.json:Zone.Identifier:Zone.Identifier
│       ├── loader_like.json
│       ├── loader_like.json:Zone.Identifier
│       ├── loader_like.json:Zone.Identifier:Zone.Identifier
│       ├── ransomware_like.json
│       ├── ransomware_like.json:Zone.Identifier
│       └── ransomware_like.json:Zone.Identifier:Zone.Identifier
├── malwarelens_ai
│   ├── analyze_sample.py
│   ├── analyze_sample.py:Zone.Identifier
│   ├── analyze_sample.py:Zone.Identifier:Zone.Identifier
│   ├── cluster_campaigns.py
│   ├── cluster_campaigns.py:Zone.Identifier
│   ├── cluster_campaigns.py:Zone.Identifier:Zone.Identifier
│   ├── config.py
│   ├── config.py:Zone.Identifier
│   ├── config.py:Zone.Identifier:Zone.Identifier
│   ├── feature_extract_safe.py
│   ├── feature_extract_safe.py:Zone.Identifier
│   ├── feature_extract_safe.py:Zone.Identifier:Zone.Identifier
│   ├── generate_rules.py
│   ├── generate_rules.py:Zone.Identifier
│   ├── generate_rules.py:Zone.Identifier:Zone.Identifier
│   ├── __init__.py
│   ├── __init__.py:Zone.Identifier
│   ├── __init__.py:Zone.Identifier:Zone.Identifier
│   ├── __pycache__
│   │   ├── analyze_sample.cpython-313.pyc
│   │   ├── analyze_sample.cpython-313.pyc:Zone.Identifier
│   │   ├── analyze_sample.cpython-313.pyc:Zone.Identifier:Zone.Identifier
│   │   ├── cluster_campaigns.cpython-313.pyc
│   │   ├── cluster_campaigns.cpython-313.pyc:Zone.Identifier
│   │   ├── cluster_campaigns.cpython-313.pyc:Zone.Identifier:Zone.Identifier
│   │   ├── config.cpython-313.pyc
│   │   ├── config.cpython-313.pyc:Zone.Identifier
│   │   ├── config.cpython-313.pyc:Zone.Identifier:Zone.Identifier
│   │   ├── feature_extract_safe.cpython-313.pyc
│   │   ├── feature_extract_safe.cpython-313.pyc:Zone.Identifier
│   │   ├── feature_extract_safe.cpython-313.pyc:Zone.Identifier:Zone.Identifier
│   │   ├── generate_rules.cpython-313.pyc
│   │   ├── generate_rules.cpython-313.pyc:Zone.Identifier
│   │   ├── generate_rules.cpython-313.pyc:Zone.Identifier:Zone.Identifier
│   │   ├── __init__.cpython-313.pyc
│   │   ├── __init__.cpython-313.pyc:Zone.Identifier
│   │   ├── __init__.cpython-313.pyc:Zone.Identifier:Zone.Identifier
│   │   ├── synthetic_corpus.cpython-313.pyc
│   │   ├── synthetic_corpus.cpython-313.pyc:Zone.Identifier
│   │   ├── synthetic_corpus.cpython-313.pyc:Zone.Identifier:Zone.Identifier
│   │   ├── train_models.cpython-313.pyc
│   │   ├── train_models.cpython-313.pyc:Zone.Identifier
│   │   └── train_models.cpython-313.pyc:Zone.Identifier:Zone.Identifier
│   ├── synthetic_corpus.py
│   ├── synthetic_corpus.py:Zone.Identifier
│   ├── synthetic_corpus.py:Zone.Identifier:Zone.Identifier
│   ├── train_models.py
│   ├── train_models.py:Zone.Identifier
│   └── train_models.py:Zone.Identifier:Zone.Identifier
├── models
│   ├── deep_mlp.joblib
│   ├── deep_mlp.joblib:Zone.Identifier
│   ├── deep_mlp.joblib:Zone.Identifier:Zone.Identifier
│   ├── static_rf.joblib
│   ├── static_rf.joblib:Zone.Identifier
│   └── static_rf.joblib:Zone.Identifier:Zone.Identifier
├── notebooks
│   ├── README_Notebooks.md
│   ├── README_Notebooks.md:Zone.Identifier
│   └── README_Notebooks.md:Zone.Identifier:Zone.Identifier
├── reports
│   ├── analysis_report.json
│   ├── analysis_report.json:Zone.Identifier
│   ├── analysis_report.json:Zone.Identifier:Zone.Identifier
│   ├── analysis_report.md
│   ├── analysis_report.md:Zone.Identifier
│   ├── analysis_report.md:Zone.Identifier:Zone.Identifier
│   ├── campaign_clusters.png
│   ├── campaign_clusters.png:Zone.Identifier
│   ├── campaign_clusters.png:Zone.Identifier:Zone.Identifier
│   ├── cluster_assignments.csv
│   ├── cluster_assignments.csv:Zone.Identifier
│   ├── cluster_assignments.csv:Zone.Identifier:Zone.Identifier
│   ├── cluster_summary.json
│   ├── cluster_summary.json:Zone.Identifier
│   ├── cluster_summary.json:Zone.Identifier:Zone.Identifier
│   ├── confusion_matrix.png
│   ├── confusion_matrix.png:Zone.Identifier
│   ├── confusion_matrix.png:Zone.Identifier:Zone.Identifier
│   ├── metrics.json
│   ├── metrics.json:Zone.Identifier
│   └── metrics.json:Zone.Identifier:Zone.Identifier
├── rules
│   ├── ai_assisted_triage.yar
│   ├── ai_assisted_triage.yar:Zone.Identifier
│   └── ai_assisted_triage.yar:Zone.Identifier:Zone.Identifier
├── scripts
│   ├── __pycache__
│   │   ├── run_full_demo.cpython-313.pyc
│   │   ├── run_full_demo.cpython-313.pyc:Zone.Identifier
│   │   └── run_full_demo.cpython-313.pyc:Zone.Identifier:Zone.Identifier
│   ├── run_full_demo.py
│   ├── run_full_demo.py:Zone.Identifier
│   ├── run_full_demo.py:Zone.Identifier:Zone.Identifier
│   ├── run_full_demo.sh
│   ├── run_full_demo.sh:Zone.Identifier
│   └── run_full_demo.sh:Zone.Identifier:Zone.Identifier
├──LICENSE
├── README.md
└── requirements.txt

Setup

git clone https://github.com/Surajit-Samanta/MalwareLens_AI.git
cd MalwareLens_AI
python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

# YARA is a native binary, not a pip package — install it separately:
sudo apt install -y yara      # Debian/Ubuntu

Usage

# 1. Generate the synthetic corpus + safe_samples/*.json
python -m malwarelens_ai.synthetic_corpus

# 2. Train both models
python -m malwarelens_ai.train_models

# 3. Analyze a single sample
python -m malwarelens_ai.analyze_sample data/safe_samples/infostealer_like.json

# 4. Cluster into likely campaigns
python -m malwarelens_ai.cluster_campaigns

# 5. Generate a draft YARA rule and check it
python -m malwarelens_ai.generate_rules data/safe_samples/ransomware_like.json
yara rules/ai_assisted_triage.yar data/safe_samples/ransomware_like.json

# Or run steps 1-5 in one shot:
python scripts/run_full_demo.py

Advanced extension: cross-checking against real EMBER data

# Get EMBER2018 raw features: https://github.com/elastic/ember
# Get EMBER2024 (via thrember): https://github.com/FutureComputing4AI/EMBER2024/

# Convert one EMBER .jsonl shard into this project's feature schema
python -m malwarelens_ai.ember_corpus data/ember2018/train_features_0.jsonl 20000

# Train on it (back up your synthetic-trained models/metrics first if you want both)
python -c "from malwarelens_ai.train_models import train; train(csv_path='data/ember_features.csv')"

# Or batch-analyze many EMBER records directly against a trained model
python -m malwarelens_ai.analyze_jsonl data/ember_2017_2/test_features.jsonl --limit 200

Using your own real files

python -m malwarelens_ai.custom_corpus build data/my_malicious_like/ --label 1
python -m malwarelens_ai.custom_corpus build data/my_benign_like/ --label 0 --append
python -c "from malwarelens_ai.train_models import train; train(csv_path='data/my_dataset.csv')"

Known limitations

  • Synthetic data is trivially separable. Each family is one fixed base feature vector plus Gaussian noise, so both models reach ~1.0 accuracy/F1/AUC on the synthetic split — that's a separability artifact of the generator, not evidence of real-world detection strength.
  • EMBER feature-parity gap. Only 8 of the 22 project features have a genuine equivalent in EMBER's raw schema; the remaining 14 (ip_count, base64_blob_count, powershell_hint_count, anti_analysis_hint_count, and several API-category counts) default to 0 for EMBER-derived records.
  • MBC_MAP only covers 5 hardcoded feature categories — any malicious behavior outside those keys is never surfaced.
  • analyze_sample.py's strings_of_interest assumes obj['strings'] is a list. If fed a raw EMBER record (where strings is a dict of statistics), it silently returns dictionary key names instead of real string content — use ember_corpus.py / analyze_jsonl.py to convert EMBER records first.
  • sha256 is only genuine for real files. Synthetic samples use a placeholder string (synthetic-{family}-0001), not an actual hash.
  • Non-deterministic reruns. Models are retrained fresh each run; without persisting a fixed model artifact, the same sample's reported probability can shift slightly between runs.

License

MIT — see LICENSE.

About

Static-analysis malware triage workbench (RF + MLP) with MBC/ATT&CK mapping, campaign clustering, and AI-assisted YARA rule generation — built for AICS-108

Topics

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages