Анализ зависимостей бинарных файлов на основе ML

от автора

Всем привет! ? ? ? Мы стажеры-разработчики Тинькофф: Влад, Паша и Илья. В проекте по стажировкам в ИБ Summer of Code под руководством Ромы Лебедя мы реализовали анализатор бинарного кода на основе ML-подходов — Binary SCA. Наш проект совмещает две предметные области — информационную безопасность и ML, поэтому мы разделили статью на несколько частей. 

В этой статье подробно расскажем о ML-стороне проекта: проведенные исследования, сложности, с которыми столкнулись в ходе работы, какой результат получили. В этой части делимся опытом использования Rizin и Milvus. Добро пожаловать! 

Использование ML в SCA

Уровень абстракции в процессах разработки непрерывно повышается за счет большого количества готовых компонентов. С одной стороны, это приводит к ускорению реализации запроса, с другой — усложняет анализ на наличие проблем безопасности. Программисты и компании постоянно ищут инструменты, которые помогут им управлять зависимостями и гарантировать безопасность своего кода. В такие моменты на помощь приходит SCA, о котором рассказывал Влад: 

Насколько использование ML в SCA может быть оправданно? Рассмотрим случай, когда никакой метаинформации о зависимостях нет или ей нельзя доверять, но есть собранный бинарный файл. Это может быть, например, полученный от стороннего разработчика компонент. 

Большинство существующих решений анализа программ на зависимости заточены как раз на наличие вспомогательной информации помимо бинарного файла, а рассматриваемый нами сценарий хоть и редкий, но самый сложный и в то же время вариативный на возможные пути решения. На помощь могут прийти способы сравнения дизассемблированного кода, разговор о которых мы вывели в отдельную статью.

Основные подходы для сравнения опираются на константы, эвристики и

Hidden text

Можно вспомнить гипотезу Коллатца, ведь для этой короткой программы люди до сих пор не знают, эквивалентна ли она простому std::cout << 1; Задача проверки двух алгоритмов на эквивалентность алгоритмически неразрешима, но все же в нашем случае есть принципиальное ограничение — две программы на входе получены из некоторой начальной программы путем преобразований, гарантирующих алгоритмическую эквивалентность. 

Люди при сравнении двух программ руководствуются в том числе и семантикой. И здесь на помощь приходит NLP, ведь именно этот раздел связан с выделением смысла из текстов. К сожалению, на тему использования ML именно для композиционного анализа практически нет исследований или работающих проектов. В первую очередь из-за того, что тема требует знаний как в ML, так и в ИБ, а множества людей, вовлеченных в каждую из областей, довольно слабо пересекаются. Кроме того, может возникнуть вопрос, насколько NLP-модели здесь применимы. Ведь они обучаются на человеческой речи и работают с ней, а мы хотим применить их к программам, да еще и на машинном коде. 

Подходы для поиска бинарных клонов

Хоть и нет проверенного ML-решения именно для компонентного анализа, подходы для поиска бинарных клонов кода и оценки сходства кода с помощью ML-моделей все же есть. В нашей работе мы опирались на опыт двух архитектур — Asm2Vec и VulHawk. 

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

Архитектура Asm2Vec

Архитектура Asm2Vec

VulHawk предлагает более сложную и актуальную архитектуру, использующую трансформер RoBERTa. В такой архитектуре ассемблерный код подготавливается перед обработкой модели: адреса, константы и имена заменяются новыми спец-токенами. За большими подробностями рекомендуем обратиться к оригинальным статьям.

Архитектура VulHawk

Архитектура VulHawk

Цель нашего проекта — разработать black-box-анализатор, который будет принимать пользовательский Jar-файл и выводить SBOM в формате CycloneDX. SBOM содержит информацию о библиотеках и их версиях, используемых в приложении.

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

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

  • Возможность быстро пополнять базу известных зависимостей. Ежедневно база пополняется новыми библиотеками и их версиями. Мы должны гибко адаптироваться под эти изменения, чтобы предсказывать все новые и новые зависимости.

  • Точное определение набора зависимостей и их версий в собранном приложений без метаданных. Мы должны корректно находить библиотеки-кандидаты с максимальным количеством TP и минимальным числом FP.

Asm2Vec и VulHawk не смогут полностью решить нашу задачу: выделить из собранного проекта его компоненты. И с ходу непонятно, как локализовать потенциальные зависимости в коде и возможно ли это сделать как-то быстро и оптимально.

VulHawk сильно завязан на supervised-подходе, а Asm2Vec хоть и предлагает некоторый metric-learning, но использует в себе устаревший Word2Vec. Мы попытались совместить некоторые идеи из этих решений. Будем строить эмбеддинги для составных частей библиотек, благо в Java собранные проекты хорошо делятся на такие части.

Собранный проект на Java — это jar-файл, который, по сути, является zip-архивом. Внутри лежат бинарные .class-файлы, которые соответствуют классам исходной программы. На стартовом этапе исследования мы решили получать эмбеддинги для классов, но позже будем строить эмбеддинги и для методов.

Концепция SCAML

Сначала нам нужно было что-то, с чем можно сравниться. Как уже было отмечено, у нас не supervised-схема, но способ проверить приемлемость точности нашей модели должен быть. Так что мы придумали не ML-подход, который учитывает довольно мощную информацию из jar-файлов — константы. Все дело в том, что в константах jar-файлов без обфускации содержатся и полные имена методов и классов, так что точность подхода ожидается высокой.

Идея алгоритма в том, чтобы в качестве библиотек, входящих в проект, брать те, константы которых полностью содержатся во множестве констант проекта. Такой подход гарантирует отсутствие FN, то есть предсказание будет содержать все истинные зависимости. Но никаких гарантий относительно FP быть не может: например, константы библиотеки могут покрываться объединением констант других библиотек, тогда получим ложную библиотеку в предсказании. Но можно положиться на то, что среди констант есть названия методов и классов: такие специфичные константы могут помочь различать между собой разные библиотеки, но не всегда разные версии. 

Реализовали подход с помощью Postgres базы данных: завели две таблички — под константы без повторений и под библиотеки, а также таблицу связи. Предварительно фильтровали константы: обрезали с конца по максимальной длине и не брали слишком мелкие, избавлялись от мусора. Тогда предсказание делается в один запрос: нужно по входным константам найти те библиотеки, константы которых полностью включаются во входное множество. Кроме того, можем получить неплохую асимптотику, если использовать индексы на столбцах таблицы, это O(n\log m), где m — количество записей в таблице, а n — количество входных констант. 

Архитектура нашего проекта: определяем, какие библиотеки хотим выявить, потом части этих библиотек векторизируем алгоритмом, сближающим похожий по семантике код и отдаляющим разный. В итоге в векторном хранилище будут лежать векторы классов известных библиотек и информация о том, из какой библиотеки эти векторы. Все это — preprocessing, флоу данных показан красным на рисунке

Архитектура нашего проекта: определяем, какие библиотеки хотим выявить, потом части этих библиотек векторизируем алгоритмом, сближающим похожий по семантике код и отдаляющим разный. В итоге в векторном хранилище будут лежать векторы классов известных библиотек и информация о том, из какой библиотеки эти векторы. Все это — preprocessing, флоу данных показан красным на рисунке

Если мы хотим выявить из входного файла зависимости, то нам нужно опять векторизовать классы входного файла. Далее использовать блок Predition для поиска зависимости в проекте, флоу данных в этом случае обозначен синим на рисунке. Prediction обращается в хранилище и на основе

Hidden text

В качестве модели-трансформера мы выбрали RoBERTa, потому что она является улучшенной версией rizin — open-source-инструмент и фреймворк для реверс-инжиниринга. Мы выбрали его, потому что:

  • у него открытый исходный код, поэтому он всегда будет доступен в России;

  • он поддерживает разные архитектуры и форматы файлов (не только Jar), что позволит в будущем облегчить анализ бинарных файлов для других платформ;

  • у него удобный API, для получения Control Flow Graph нужно всего несколько простых команд:

icm — получить список методов в классе  s <addr> — переместиться на начало нужного метода  af — проанализировать метод  agf — вывести CFG

Мы начали пользоваться rizin’ом и поняли, что поддержка jvm в нем несовершенна:

  • некорректно обрабатывались опкоды tableswitch и lookupswitch, в результате некорректно дизассемблировались switch-case-конструкции; 

  • rizin не понимал, что throw — это конец пути исполнения кода, и пытался дизассемблировать опкоды после него;

  • в целом в rizin не учитывается, что в jvm нельзя прыгать за пределы метода, из-за этого неправильно распаршенные опкоды часто дизассемблировались как прыжки в другие части class-файла;

  • и еще пара случаев, когда rizin некорректно парсил опкоды: иногда считал, что case-1 — это невалидный опкод или некорректно обрабатывал wide инструкции jvm.

Hidden text
Код, дизасемблированный rizin до патча. Видно, как некорректно парсится опкод tableswitch

Код, дизасемблированный rizin до патча. Видно, как некорректно парсится опкод tableswitch
Код, дизасемблированный rizin после патча

Код, дизасемблированный rizin после патча

Мы потратили пару недель на доработку инструмента. Возможно, в будущем откроем pull-реквест в апстрим.

Мы научились дизассемблировать библиотеки для получения фич, на их основе с помощью RoBERTa мы получаем вектора и складываем в некоторое хранилище. Пусть нам на вход пришел собранный проект для анализа, содержащий известную нам библиотеку X. Тогда заметим следующее:

Если этот проект при сборке не был обфусцирован, то множество его векторов включает в себя векторы X — красные точки. Если же проект как-либо обфусцирован, то мы ожидаем, что векторы, отвечающие классам библиотеки X, будут близки к оригинальным версиям

Если этот проект при сборке не был обфусцирован, то множество его векторов включает в себя векторы X — красные точки. Если же проект как-либо обфусцирован, то мы ожидаем, что векторы, отвечающие классам библиотеки X, будут близки к оригинальным версиям

Алгоритм получился такой: 

— для каждого вектора входного файла находим n_neighbours ближайших, лежащих в некоторой \varepsilon-окрестности;

— получаем некоторое множество Y векторов из базы данных, для которых мы знаем, из какой библиотеки они пришли;

— находим те библиотеки, классы которых полностью включаются в извлеченное множество Y. Их и будем выдавать в качестве предсказания, ведь все их классы оказались рядом с классами входного файла. 

Остается пара неясных моментов:

  1. Какую базу данных для векторов использовать, чтобы иметь возможность делать быстрый поиск ближайших и пополнять новыми зависимостями?

  1. Насколько эта схема вообще оправданна? Ведь она основывается на предположении, что при обфускации эмбеддинг класса окажется рядом с исходным, и нам нужен какой-то способ это проверить.

Для хранения векторов мы решили использовать Milvus — open source векторную базу данных, предоставляющую все необходимые для нашей задачи функции. Чтобы обеспечить быстрый поиск ближайших векторов, Milvus предлагает различные Схематичный пример работы IVF_FLAT: q — точка запроса, c_0 и c_1 — ближайшие центры кластеров, тогда поиск будет проводиться среди точек v_0,...v_6

Схематичный пример работы IVF_FLAT: q — точка запроса, c_0 и c_1 — ближайшие центры кластеров, тогда поиск будет проводиться среди точек v_0,…v_6

Milvus в реализации ANNS использует низкоуровневые оптимизации: например, векторы запроса кладут в L3-кеш процессора, чтобы уменьшить количество обращений в ОЗУ, реализуют поиск на нескольких ядрах с помощью OpenMP. 

Векторы запроса — q0,...,q_m — делятся на блоки, каждый из которых кладется в L3-кеш. Каждый из t потоков считает расстояние между векторами из базы v_0,...,v_{b-1} и текущим блоком в L3 — результаты кладутся в t соответствующих куч, итоговый результат получается слиянием куч

Векторы запроса — q0,…,q_m — делятся на блоки, каждый из которых кладется в L3-кеш. Каждый из t потоков считает расстояние между векторами из базы v_0,…,v_{b-1} и текущим блоком в L3 — результаты кладутся в t соответствующих куч, итоговый результат получается слиянием куч

Будем складывать в milvus вектора с информацией, из какой библиотеки они пришли.

Тестирование и проверки, что подход рабочий

Мы искали способ убедиться, что наш алгоритм анализа действительно устойчив к обфускации. Для этого выделили датасет на 3000 maven-библиотек, получили векторы для них и положили на mivlus. На вход модели давали каждую из библиотек и ожидали на выходе получить ее же. Считали совпадение

Классические метрики многоклассовой классификации: precision показывает долю верно предсказанных зависимостей в предсказании, recall — долю верно предсказанных зависимостей среди истинных

Классические метрики многоклассовой классификации: precision показывает долю верно предсказанных зависимостей в предсказании, recall — долю верно предсказанных зависимостей среди истинных

Для проверки возможностей детектировать выбранным методом обфусцированные библиотеки решили подавать те же самые библиотеки, но обфусцированные с помощью Proguard. Хоть и даем на входе одну библиотеку, мы все еще имеем дело с многоклассовой классификацией, ведь на выходе можем получить несколько библиотек. Так что предлагается оценивать результаты поэлементно, для каждой библиотеки считая precision ее предсказания.

При первых попытках испытать модель выяснилось, что из-за повторов векторов часто возникали ситуации, что эмбеддинг истинной библиотеки не находился. К примеру, 10 векторов совпали с истинным, тогда мы можем не увидеть истинный. Пришлось добавить эвристики поиска, и первое, что пришло в голову, — учитывать только уникальные векторы. Вот результаты для n_neighbours=7:

Все еще много случаев с precision=0, так что мы продолжили исследовать частные примеры и выяснили, что разные версии одной и той же библиотеки часто оказываются неразличимы: множества их векторов оказываются достаточно близки или же вообще совпадают. Поэтому мы поменяли определение уникальности — теперь считаем уникальным тот класс, который встречается в пределах разных версий одной библиотеки. В итоге получили метрики:

Как видим, на библиотеках без обфускации все метрики выросли — если брать совпадение с точностью до библиотеки, то здесь средний preсision вообще 0.96, очень хороший результат. Среднее метрик при обфускации просело, но на самом деле это из-за учащения FP исходов, а не из-за случаев с precision=0 — их количество практически не изменилось, а где-то даже упало.

Стоит отметить, что в результатах выше была использована предобученная модель. Основная проблема с дообучением здесь в разметке данных: какие элементы мы можем считать похожими, а какие нет. Мы пробовали дообучать на триплетах, считая похожими изменившиеся классы в разных версиях одной библиотеки, но это не дало результата: метрики на этом же датасете стали даже хуже. Но поисследовать эмбеддинги можно будет позже, пока что нужно сконцентрироваться на том, чтобы протестировать модель на большом датасете.

Что касается бейзлайн-подхода, то также проверили аналогичные метрики на нашем датасете. Только метрик для обфусцированных версий на рисунке нет — там нулевая точность, потому что имена методов и классов затерты.

Выводы и несколько полезных ссылок

Использование полученных эмбеддингов имеет ряд допущений и, как следствие, набор ограничений. Если не брать продвинутые приемы обфускации кода, использование эмбеддингов кода для выявления зависимостей можно развивать. Основное преимущество такого подхода — относительно быстрый поиск по базе индексированного кода, быстрое добавление новых пакетов в индекс. 

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

На прощание несколько ссылок для тех, кто хочет больше погрузиться в тему:

  1. Статья с задачей о схожести бинарного кода на sciencedirect.com. Предлагают использовать гибридную сетку — CNN + LSTM. Анализ разделяется на три этапа. Во-первых, используется IDA Pro для статического анализа и дизассемблирования. На втором этапе используется моделька Word2Vec + сетка, чтобы классифицировать архитектуру процессора и уровни оптимизации компилятора. После этого нейронная сеть Siamese NN детектирует схожесть.

  2. В статье на ndss-symposium.org предлагают подход к решению задачи поиска похожих блоков кода в двух входных бинарных файлов. Строят эмбеддинги блоков с помощью Word2Vec, но дополнительно вводятся веса для инструкций — их TF-IDF. Далее приближенно решается задача о назначениях с помощью жадного алгоритма.

  3. В статье на ieeexplore.ieee.org предлагают решение задачи поиска клонов кода на основе графовой информации: CFG, AST, https://habr.com/ru/articles/796919/


Комментарии

Добавить комментарий

Ваш адрес email не будет опубликован. Обязательные поля помечены *