Malware Classifier: Статический анализ PE-файлов

Дипломный проект — система машинного обучения для обнаружения вредоносного ПО путем статического анализа PE-файлов (Portable Executable) Windows. Проект включает в себя полный цикл: от извлечения признаков до предсказания с использованием различных моделей.

Важно: Проект находится в активной разработке, структура и файлы могут часто обновляться.

Возможности

• Анализ отдельных файлов или целых директорий (рекурсивно).
• Режим "сканера" — автоматическая сортировка файлов на "чистые" (benign) и "вредоносные" (malware).
• Сохранение отчётов о результатах проверки в текстовый файл.
• Поддержка нескольких моделей: на данный момент доступны SGD и CatBoost классификаторы.

Структура проекта (для инференса)

Для работы с предсказаниями вам потребуется следующая структура файлов и папок:

malware_predict/                  # Корневая папка проекта (может называться иначе)
│
├── _3_all_predicts/               # Скрипты для предсказания
│   ├── sgd/
│   │   └── 7.predict.py           # Скрипт для моделей SGD
│   └── cats/
│       └── 9.predict_catboost.py  # Скрипт для модели CatBoost
│
├── models/                         # Обученные модели и словари
│   ├── sgd/
│   │   ├── l1_alpha_10-4/          # Пример папки с моделью SGD (L1 рег.)
│   │   │   ├── malware_classifier.pkl
│   │   │   └── vocabulary.pkl
│   │   ├── l2_alpha_10-3/          # Пример папки с моделью SGD (L2 рег.)
│   │   │   ├── malware_classifier.pkl
│   │   │   └── vocabulary.pkl
│   │   └── l2_tfidf_alpha_10-3/    # Пример папки с моделью SGD (L2 рег. TF IDF)
│   │       ├── malware_classifier.pkl
│   │       ├── tfidf_transformer.pkl
│   │       └── vocabulary.pkl
│   └── cats/
│       └── i_626/                  # Пример папки с моделью CatBoost
│           ├── catboost_malware_classifier.cbm
│           ├── vocabulary_dict.pkl  # Основной словарь для CatBoost
│           ├── vocabulary.pkl       # Дополнительный файл (может не использоваться)
│           └── training_metadata.json
│
└── classes/                         # Вспомогательные классы (необходимы для работы)
    └── pe_analyzer.py               # Класс для анализа PE-файлов

Быстрый старт

Предварительные требования

• Windows 10/11
• Python 3.12 (Последняя версия модели Cat работает только на 3.12)
• Git (для клонирования репозитория)
• Git LFS (Git Large File Storage) — обязателен, так как модели и словари имеют большой размер.

Установка

1. Клонируйте репозиторий (с поддержкой Git LFS):

   # Убедитесь, что Git LFS установлен (git lfs install)
   git clone https://github.com/jax225/malware_predict/
   cd malware_predict

   Альтернативный вариант: Если вы не используете Git, скачайте и распакуйте проект как ZIP-архив. В этом случае вам потребуется вручную создать структуру папок (classes, _3_all_predicts/sgd, и т.д.) и скопировать в неё необходимые файлы.

2. Установите зависимости:

   pip install pefile numpy scipy scikit-learn py7zr rarfile catboost

Запуск предсказания

Способ запуска зависит от выбранной модели.

Для моделей SGD (Стохастический градиентный спуск)

1. Перейдите в папку со скриптом:

   # Перейдите в основную папку директории и запускайте предсказания из нее. 
   cd malware_predict

2. Выполните скрипт, указав пути к модели, словарю и целевому файлу/папке. Пути к classes определяются автоматически относительно скрипта.

   python .\_3_all_predicts\sgd\7.predict.py --model ".\models\sgd\l1_alpha_10-4\malware_classifier.pkl" --vocab ".\models\sgd\l1_alpha_10-4\vocabulary.pkl" --test

   python .\_3_all_predicts\sgd\7.predict.py --model ".\models\sgd\l2_alpha_10-3\malware_classifier.pkl" --vocab ".\models\sgd\l2_alpha_10-3\vocabulary.pkl" --test

   python .\_3_all_predicts\sgd\7.predict.py --model ".\models\sgd\l2_tfidf_alpha_10-3\malware_classifier.pkl" --vocab ".\models\sgd\l2_tfidf_alpha_10-3\vocabulary.pkl" --tfidf .\models\sgd\l2_tfidf_alpha_10-3\tfidf_transformer.pkl  --test

   # Для модели CatBoost используется 9.predict_catboost.py
   python .\_3_all_predicts\cats\9.predict_catboost.py --model ".\models\cats\i_626\catboost_malware_classifier.cbm" --vocab ".\models\cats\i_626\vocabulary_dict.pkl" --test

Справка по аргументам

Для получения полного списка аргументов используйте ключ -h:

python 7.predict.py -h
python 9.predict_catboost.py -h

Основные аргументы:

• --model : Путь к файлу обученной модели.
• --vocab : Путь к файлу словаря (vocabulary).
• --tfidf : Путь к файлу перобразование (используется только для TF-IDF модели)
• --file : Путь к анализируемому файлу.
• --dir : Путь к анализируемой папке (будет просканирована рекурсивно).
• --sort : Включить режим сортировки файлов по папкам malware и benign.
• --report: Сохранить результаты проверки в файл scan_report.txt.

Как это работает

0. Скрипт не открывает файл для проверки, а извлекает библиотеки. Нет реального открытия файла.
1. Анализ PE-файла: Скрипт (pe_analyzer.py) извлекает из файла статические признаки: информацию из заголовков, список секций, импортируемые функции, энтропию и сырые строки.
2. Векторизация: Извлеченные строки преобразуются в числовой вектор на основе предоставленного vocabulary.pkl.
3. Предсказание: Вектор признаков подается на вход загруженной модели (malware_classifier.pkl или .cbm).
4. Результат: Модель возвращает вердикт: вредоносный файл (Malware) или чистый (Benign).

Модели и производительность

В проекте исследуются различные подходы к классификации. На данный момент реализованы:

• SGD Classifier: Две линейные модели с L1 и L2 регуляризацией. Легкие и быстрые.
• SGD Classifier: Линейная модель с L2 регуляризацией. с TF-IDF модификацией. (модель в первую очередь смотрит на уникальные признаки)
• CatBoost: Градиентный бустинг, который часто показывает более высокую точность на сложных зависимостях.

Целевые метрики: F1-Score > 0.97 на тестовых выборках.

Планы по развитию

• Реализация методов сглаживания (сглаживание Лапласа) для борьбы с переобучением.
• Глубокое исследование влияния L1 и L2 регуляризации.
• Создание ансамбля моделей: LightGBM, нейронные сети и стекинг для повышения качества детектирования.

Полная документация

Подробное описание процесса обучения, подготовки данных и экспериментальных результатов можно найти в файле: Процесс обучения модели и предсказания

Полный документ по запуску модели приведен в файле doc/Runbook.md

---

Проект разработан в рамках дипломной работы. Код и модели доступны для некоммерческого использования.