Под капотом у MOLOT’а: обзор ML-модели для обнаружения вредоносного кода

Всем привет! На связи Максим Митрофанов, руководитель ML-команды Application Security в Positive Technologies.

Недавно мы выпустили нейросеть MOLOT и описали архитектуру решения в техрепорте на arXiv. Эта статья в каком-то смысле его TLDR в свободном изложении. Скажу сразу, наша модель — не LLM, хотя к трансформерам имеет самое прямое отношение. MOLOT построен на базе BERT-энкодера (не modern, не neo, а самого ванильного) и умеет обнаруживать вредоносный код по цепочкам вызовов не хуже Mythos от Anthropic (как минимум не боится промпт-инъекций ^[1]).

Malicious code (вредоносный код) — это фрагменты кода, которые намеренно добавляют в проект для скрытого выполнения опасных действий: кражи данных, получения удаленного доступа, обхода защиты и других.

Для тех, у кого много сил, времени и желания, мы открыли бенчмарк, на котором были подсчитаны метрики. Ссылку на техрепорт также прикладываю:

• MOLOT System Card ^[2]

• Open Malicious Code Benchmark ^[3]

Ниже разберем самые важные детали пилотного релиза ML-модуля.

Не все модели LLM, не всем нужен GPU

Сегодня принято под любую проблему брать LLM, даже там, где может хватить моделей попроще. Как я уже сказал, MOLOT — это энкодер. Такой выбор обусловлен не только качественными метриками, но и нефункциональными. Мы изначально целились в локальную работу, чтобы хватало производительности и пропускной способности памяти ^[4] даже процессора. Вот системные требования PT Application Inspector, модулем которого и является MOLOT.

Производительность, которую показывает MOLOT, при сопоставимых размерах модели пока недостижима на больших языковых моделях сопоставимого размера. Поэтому языковые модели 1b мы отмели сразу.

Еще одно важное отличие от LLM: encoder-архитектура позволяет избежать галлюцинаций. Это достигается за счет другой постановки задачи при обучении ^[5]: вместо предсказания следующего токена (основная причина галлюцинаций) модель смотрит на всю последовательность и принимает финальное решение, без вероятностных механизмов декодеров.

Почему мы выбрали BERT

В какой-то момент разработки мы задались вопросом, стоит ли изменить что-то в самой модели. Сперва мы взялись за доменно-специфичные варианты, обратив внимание ^[6] на ИБ и кодинг-специфичные базовые модели.

Метрики токенизации менялись в пределах погрешности, а конечное качество файнтюнинга и вовсе не различалось.

Эксперименты по расширению контекста также оказались противоречивыми. Мы взяли ModernBERT как одну из последних моделей с расширенным контекстом, ограничив его до 4096 токенов.

Можно заметить, что с увеличением контекста итоговое качество не изменилось. А более точечный анализ ошибок показал, что новые детекты, охватившие больший контекст, сменились пропусками вредоносного кода в начале цепочки вызовов. При этом важно отметить, что производительность ModernBERT падала обратно пропорционально увеличению контекста.

По совокупности результатов мы приняли решение остаться на ванильном bert-base-uncased и вложить усилия в качественную подготовку данных и интерпретируемость результатов.

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

Работу MOLOT можно разбить на пять шагов:

1. Достаем из кода граф вызовов и переменные. Анализ исходного кода ИИ-моделями возможен, но неэффективен: код может быть обфусцирован, содержать неинформативные комментарии или иметь сложную логику ^[7] выполнения, которую будет неправильно читать сверху вниз.

2. Разбиваем граф на цепочки вызовов для каждого файла. В первую очередь это важно для оптимизации контекста модели. Пофайловые вердикты обеспечивают соответствие требованиям, предъявляемым к SAST-инструментам. В контексте этих требований, помимо вердикта по проекту, важна локализация проблемы. Кроме того, нам встречались хакерские кампании, в которых использовался язык, отличный от основного языка проекта.

3. Переходим от цепочки вызовов к цепочке активностей. Этот подход мы взяли из решения Cerebro ^[8]. В техрепорте ^[3] подробно рассказываем, как устранили критические недостатки оригинального решения и развили подход для более качественного и стабильного статического анализа кода .

4. Классифицируем. Модель получает на вход цепочку активностей и выдает вердикт о вредоносности файла проекта.

5. Интерпретируем результат. Связность всех этапов работы ML-модуля позволила развернуть важные токены цепочки активности обратно в строчки исходного кода. Мы подобрали оптимальные параметры для SHAP-интерпретации, которые дали 81% согласованности по строкам, выбранным экспертами в процессе разметки вредоносного кода.

Одна из ключевых проблем в работе нейронных сетей — трудность интерпретации результатов. Поскольку срабатывания SAST разбирают специалисты по безопасной разработке, нам было важно обосновывать выданное моделью решение. Результата на скриншоте ниже мы достигли за счет кастомизации SHAP-анализа ^[9] и обеспечения связности цепочки активностей с графом и кодом — это позволило показать пользователям строки, на которые модель обратила внимание.

OMCBench и результаты

Для нас важно соответствовать высоким стандартам качества и оправдывать доверие сообщества и пользователей. Поэтому мы опубликовали Open Malicious-Code Benchmark ^[3] (OMCBench), который содержит 400 вредоносных и 400 легитимных пакетов из реестров PyPI, npm. Он является частью нашего закрытого бенчмарка на 4,5 тыс. примеров и объективно отражает метрики на меньшем их количестве примеров, что подтверждается доверительными интервалами, приведенными в техническом отчете ^[3]. Все метрики ниже были получены на OMCBench и могут быть перепроверены пользователями и энтузиастами сообщества безопасной разработки.

Мы сравнили MOLOT с открытыми решениями для обнаружения вредоносного кода, собрав самые популярные фреймворки безопасной разработки:

• Microsoft OSSGadget ^[10] (oss-detect-backdoor ^[11]) — запускался через Microsoft AppInspector / Application Inspector ^[12] с использованием набора правил BackdoorRules.

• bandit4mal ^[13] — форк Bandit с правилами для обнаружения вредоносного кода, только для Python; см. также связанную работу A benchmark comparison of Python malware detection approaches ^[14].

• semgrep ^[15] — использовался с правилами из набора apiiro/malicious-code-ruleset ^[16].

На всех языках, которые участвовали в сравнении (JavaScript, TypeScript, Python), MOLOT превосходит аналоги, использующие механизм сравнения по шаблону (pattern-matching). Разница по числу ложных срабатываний (Precision) доходит до 40%, по полной метрике качества (F1) — до 15%. Это еще раз подчеркивает важность анализа цепочки активностей в задаче обнаружения вредоносного кода.

Дальнейшее развитие

В первом релизе MOLOT мы поддержали три самых популярных языка для веб-разработки: TypeScript, JavaScript и Python. Мы продолжим изучать тренды и потребности ^[17] пользователей, чтобы в дальнейшем расширить список поддерживаемых языков. Мы также следим за развитием хакерских техник в части внедрения вредоносного кода. Благо мы работаем в компании в сфере ИБ и можем доливать уникальные кейсы в обучающий набор данных следующих версий.

Спасибо всем, кто помог дойти до релиза MOLOT, кто спорил, убеждал и кто поверил. Если было интересно, то подписывайтесь (даже если не было, все равно подписывайтесь) на наши каналы в Telegram, где мы проводим открытые ридинг-группы и разбираем кейсы:

• False Positive ^[18] — сообщество по ML и AI в кибербезопасности

• Positive Development Community ^[19] — сообщество по безопасной разработке

• ESCalator ^[20] — сообщество экспертов по кибербезопасности