Всем привет! ? ? ? Мы стажеры-разработчики Тинькофф: Влад, Паша и Илья. В проекте по стажировкам в ИБ 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 здесь доработан, чтобы учитывать особенности синтаксиса ассемблера. В итоге у них получилось, что близкие по функциям опкоды были близки в векторном пространстве.
VulHawk предлагает более сложную и актуальную архитектуру, использующую трансформер RoBERTa. В такой архитектуре ассемблерный код подготавливается перед обработкой модели: адреса, константы и имена заменяются новыми спец-токенами. За большими подробностями рекомендуем обратиться к оригинальным статьям.
Цель нашего проекта — разработать 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 базы данных: завели две таблички — под константы без повторений и под библиотеки, а также таблицу связи. Предварительно фильтровали константы: обрезали с конца по максимальной длине и не брали слишком мелкие, избавлялись от мусора. Тогда предсказание делается в один запрос: нужно по входным константам найти те библиотеки, константы которых полностью включаются во входное множество. Кроме того, можем получить неплохую асимптотику, если использовать индексы на столбцах таблицы, это 
, где m — количество записей в таблице, а n — количество входных констант.
Если мы хотим выявить из входного файла зависимости, то нам нужно опять векторизовать классы входного файла. Далее использовать блок Predition для поиска зависимости в проекте, флоу данных в этом случае обозначен синим на рисунке. Prediction обращается в хранилище и на основе некоторого алгоритма выдает наиболее вероятные библиотеки и их версии.
Выбор инструментов
Мы экспериментировали с разными моделями и пробовали построить Jar2Vec по аналогии с Asm2Vec, запускали BERT на константах, на сыром бинарнике и так далее. У каждой модели было свое преимущество, и мы пришли к такой концепции:
На вход принимаем JAR-файл.
Дизассемблируем с помощью Rizin и получаем CFG — Control-flow graph.
Обход по данному графу с помощью DeepWalk на методах классов и получаем «предложения» — возможные варианты исполнения метода. В этом же пункте мы убираем из кода имена вызываемых классов и адреса, как это было в VulHawk: программа должна концентрироваться именно на контексте исполнения программы и не должна переобучаться на какие-либо названия, которые могут легко измениться от сборки к сборке.
Получаем эмбеддинги методов с помощью модели-трансформера. Чтобы получить эмбеддинг класса, просто берем среднее эмбеддингов его методов.
Hidden text
В качестве модели-трансформера мы выбрали RoBERTa, потому что она является улучшенной версией модели BERT, в которой дополнительные параметры и техники обучения были оптимизированы для достижения лучшей производительности. От обычного BERT он отличается по пяти пунктам:
Обучение на большем объеме данных: RoBERTa обучается на значительно большем объеме текстовых данных, чем BERT. Больший объем данных позволяет RoBERTa улучшить свое представление языка и семантическое понимание.
Оптимизация гиперпараметров: RoBERTa использует оптимизированные гиперпараметры, которые были тщательно подобраны для достижения более высокой производительности. Некоторые гиперпараметры, такие как размер пакета обучения и количество эпох, настроены в RoBERTa для более эффективного обучения.
Модель перестала обучаться на предсказание следующего предложения, что позволяет модели лучше и эффективнее обрабатывать контекст и информацию в предложениях.
Использование дополнительных значений предварительного обучения: RoBERTa использует дополнительные значения предварительного обучения, такие как длинные последовательности, динамическое обучение и т. д. Это помогает модели получать более точное представление языка и лучше справляться со сложными задачами обработки естественного языка.
Использует byte-level BPE-токенайзер, что лучше подходит под синтаксис.
В качестве инструмента для извлечения признаков из Jar-файлов мы выбрали 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
Мы потратили пару недель на доработку инструмента. Возможно, в будущем откроем pull-реквест в апстрим.
Мы научились дизассемблировать библиотеки для получения фич, на их основе с помощью RoBERTa мы получаем вектора и складываем в некоторое хранилище. Пусть нам на вход пришел собранный проект для анализа, содержащий известную нам библиотеку X. Тогда заметим следующее:
Алгоритм получился такой:
— для каждого вектора входного файла находим n_neighbours ближайших, лежащих в некоторой
-окрестности;
— получаем некоторое множество Y векторов из базы данных, для которых мы знаем, из какой библиотеки они пришли;
— находим те библиотеки, классы которых полностью включаются в извлеченное множество Y. Их и будем выдавать в качестве предсказания, ведь все их классы оказались рядом с классами входного файла.
Остается пара неясных моментов:
Какую базу данных для векторов использовать, чтобы иметь возможность делать быстрый поиск ближайших и пополнять новыми зависимостями?
Насколько эта схема вообще оправданна? Ведь она основывается на предположении, что при обфускации эмбеддинг класса окажется рядом с исходным, и нам нужен какой-то способ это проверить.
Для хранения векторов мы решили использовать Milvus — open source векторную базу данных, предоставляющую все необходимые для нашей задачи функции. Чтобы обеспечить быстрый поиск ближайших векторов, Milvus предлагает различные ANNS-алгоритмы, которые зачастую подразумевают предварительное построение на данных некой структуры — индекса.
Разные индексы предоставляют разный трейд-офф между скоростью, используемыми ресурсами и точностью поиска. Например, индекс, значащийся в документации как IVF_FLAT, кластеризует точки датасета и производит поиск в два этапа: сначала находит nprobe ближайших кластеров, а потом линейным поиском ищет ближайшие внутри этих кластеров.
Milvus в реализации ANNS использует низкоуровневые оптимизации: например, векторы запроса кладут в L3-кеш процессора, чтобы уменьшить количество обращений в ОЗУ, реализуют поиск на нескольких ядрах с помощью OpenMP.
Будем складывать в milvus вектора с информацией, из какой библиотеки они пришли.
Тестирование и проверки, что подход рабочий
Мы искали способ убедиться, что наш алгоритм анализа действительно устойчив к обфускации. Для этого выделили датасет на 3000 maven-библиотек, получили векторы для них и положили на mivlus. На вход модели давали каждую из библиотек и ожидали на выходе получить ее же. Считали совпадение с точностью до версии и с точностью до библиотеки.
Модель может выдавать на входе несколько пакетов, по сути, это многоклассовая классификация, так что имеет смысл для каждой библиотеки считать precision и recall предсказания для нее. Мы смотрели лишь на precision — в нашем случае recall может быть либо 0, либо 1, эту информацию увидим из анализа precision.
Для проверки возможностей детектировать выбранным методом обфусцированные библиотеки решили подавать те же самые библиотеки, но обфусцированные с помощью Proguard. Хоть и даем на входе одну библиотеку, мы все еще имеем дело с многоклассовой классификацией, ведь на выходе можем получить несколько библиотек. Так что предлагается оценивать результаты поэлементно, для каждой библиотеки считая precision ее предсказания.
При первых попытках испытать модель выяснилось, что из-за повторов векторов часто возникали ситуации, что эмбеддинг истинной библиотеки не находился. К примеру, 10 векторов совпали с истинным, тогда мы можем не увидеть истинный. Пришлось добавить эвристики поиска, и первое, что пришло в голову, — учитывать только уникальные векторы. Вот результаты для n_neighbours=7:
Все еще много случаев с precision=0, так что мы продолжили исследовать частные примеры и выяснили, что разные версии одной и той же библиотеки часто оказываются неразличимы: множества их векторов оказываются достаточно близки или же вообще совпадают. Поэтому мы поменяли определение уникальности — теперь считаем уникальным тот класс, который встречается в пределах разных версий одной библиотеки. В итоге получили метрики:
Как видим, на библиотеках без обфускации все метрики выросли — если брать совпадение с точностью до библиотеки, то здесь средний preсision вообще 0.96, очень хороший результат. Среднее метрик при обфускации просело, но на самом деле это из-за учащения FP исходов, а не из-за случаев с precision=0 — их количество практически не изменилось, а где-то даже упало.
Стоит отметить, что в результатах выше была использована предобученная модель. Основная проблема с дообучением здесь в разметке данных: какие элементы мы можем считать похожими, а какие нет. Мы пробовали дообучать на триплетах, считая похожими изменившиеся классы в разных версиях одной библиотеки, но это не дало результата: метрики на этом же датасете стали даже хуже. Но поисследовать эмбеддинги можно будет позже, пока что нужно сконцентрироваться на том, чтобы протестировать модель на большом датасете.
Что касается бейзлайн-подхода, то также проверили аналогичные метрики на нашем датасете. Только метрик для обфусцированных версий на рисунке нет — там нулевая точность, потому что имена методов и классов затерты.
Выводы и несколько полезных ссылок
Использование полученных эмбеддингов имеет ряд допущений и, как следствие, набор ограничений. Если не брать продвинутые приемы обфускации кода, использование эмбеддингов кода для выявления зависимостей можно развивать. Основное преимущество такого подхода — относительно быстрый поиск по базе индексированного кода, быстрое добавление новых пакетов в индекс.
После вводных экспериментов мы приступили к раскатке модели на больший датасет размером порядка 105 тысяч библиотек. Мы столкнулись с множеством сложностей, начиная с несостоятельности нашей схемы хранения данных, заканчивая проблемами с данными и невозможностью точно оценить качество модели. Но о том, как мы преодолевали эти вызовы, расскажем в следующей статье.
На прощание несколько ссылок для тех, кто хочет больше погрузиться в тему:
Статья с задачей о схожести бинарного кода на sciencedirect.com. Предлагают использовать гибридную сетку — CNN + LSTM. Анализ разделяется на три этапа. Во-первых, используется IDA Pro для статического анализа и дизассемблирования. На втором этапе используется моделька Word2Vec + сетка, чтобы классифицировать архитектуру процессора и уровни оптимизации компилятора. После этого нейронная сеть Siamese NN детектирует схожесть.
В статье на ndss-symposium.org предлагают подход к решению задачи поиска похожих блоков кода в двух входных бинарных файлов. Строят эмбеддинги блоков с помощью Word2Vec, но дополнительно вводятся веса для инструкций — их TF-IDF. Далее приближенно решается задача о назначениях с помощью жадного алгоритма.
В статье на ieeexplore.ieee.org предлагают решение задачи поиска клонов кода на основе графовой информации: CFG, AST, PDG.