TLDR; За счёт новой модели удалось улучшить качество распознавания истинных секретов с 0.70 до 0.90 PR AUC.
В CodeScoring мы не отстаём от трендов и активно внедряем машинное обучение в разрабатываемый продукт. В этой статье я, Антон Володченко, руководитель разработки продуктов, и Дарья Черешнюк, инженер-программист, разбираем пример использования модели для обработки результатов сканирования кода на наличие секретов.
Статья будет полезна рядовым разработчикам и специалистам в сфере информационной безопасности, которые хотят выстроить процесс разработки, исключив риски утечки конфиденциальных данных. А ещё пригодится интересующимися машинным обучением и практическими примерами его применения.
С развитием инструментов для обеспечения безопасности информационных систем всё ещё значительную роль в атаках играют утекшие пароли, API-токены и другие конфиденциальные данные. Например, исследование GitGuardian показывает, что 4.6% всех публичных репозиториев содержат секреты, а 15% авторов хотя бы раз коммитили секрет в репозиторий. С ростом популярности вайб-кодинга и использования LLM для разработки приложений эти проблемы становятся только острее, так как модели часто не заботятся о фильтрации данных в проектах и не знают, что код может сразу заливаться в публичный репозиторий. А избавиться от него там не так уж и тривиально. Это хорошо показано в статье Призраки в коммитах: как я заработал $64 000 на удаленных файлах в Git.
В борьбе с такими утечками часто используются инструменты для поиска секретов в коде, такие как Gitleaks, TruffleHog, Kingfisher. Их работа, как правило, сводится к поиску по регулярным выражениям в комбинации с другими методами, такими как оценка энтропии потенциальных секретов, попытки авторизации с токеном в сервисе, проверка наличия в защищённом хранилище вроде Vault и т.п. Эти дополнительные механизмы нужны из-за того, что обычный поиск по регулярным выражениям часто даёт большое количество ложноположительных результатов (False Positives) со всеми вытекающими последствиями в виде нехватки человеческих ресурсов на разбор, утонувшие в куче ложных срабатываний истинные секреты и т.д. Но усложнение механизмов обработки может привносить и другие варианты ложных срабатываний – ложно негативные (False Negatives). Это ситуации, когда действительно важные секреты были пропущены сканером.
Довольно простым, но при этом всё ещё актуальным примером такого поведения может являться фильтрация по низкой энтропии в Gitleaks. Рассмотрим следующий пример:
{ password := "s0m3s3cr3t" nonMatchingToken := "7lggmk1pmht76fzo" }
Здесь Gitleaks найдёт только “nonMatchingToken”. Пароль при этом будет отброшен из-за низкой энтропии.
конфигурации по умолчанию есть правило generic-api-key, в котором задан порог энтропии в 3.5. Все результаты ниже этого порога будут отсекаться. Если понизить порог энтропии, то результат уже будет полнее:
Но в общем случае при таком изменении сразу возникает проблема увеличения числа ложноположительных срабатываний. Одним из решений становится использование машинного обучения. Это стало мотивацией для создания собственного модуля поиска секретов в платформе CodeScoring, который наша команда запустила в октябре 2024 года. В первой версии модуля был использован движок Gitleaks и собственная предобученная модель с возможностью дообучения в контуре заказчика. В базовом сценарии этот подход позволяет сократить на 70% число ложноположительных срабатываний, а после дообучения их станет ещё меньше. Следующим шагом мы решили улучшить статистику за счёт использования более мощной базовой модели и дополнительной подготовки тестовых данных.
Поиск решения
Сокращение количества False Positives с точки зрения машинного обучения может быть представлено в виде задачи бинарной классификации: действительно ли найденный кандидат является секретом?
Ранее анализ кандидатов в модуле Секретов CodeScoring осуществлялся с помощью алгоритма XGBoost на признаках, описывающих текст секрета и окружающий его контекст. Среди таких признаков мы рассматривали как банальные длину и энтропию потенциального секрета, так и информацию о ключевых словах и пути к скрипту, содержащему этот секрет. Такие признаки помогали отсеять скрипты и файлы, содержащие в своём названии слова test, example и подобные, указывающие на то, что секретная информация не уйдёт дальше системы разработчика.
Так, следующего кандидата на признание его секретом:
apache/groovy/gradle/verification-metadata.xml
<trusted-key id="fc2c31fc25fede4e7ad0d18c2bfd7825a8984fbe" group="com.github.javaparser"/>\n
можно описать, например, следующими признаками:
entropy | 3.73 |
has_parentheses | FALSE |
has_brackets | FALSE |
starts_with_dollar_sign | FALSE |
has_password_word | FALSE |
has_spaces | FALSE |
has_html_tags | FALSE |
starts_with_comment | FALSE |
has_arrow | FALSE |
has_test_example_substring | FALSE |
has_programming_keywords | FALSE |
is_numerical_value | FALSE |
is_configuration_file | FALSE |
is_settings_file | FALSE |
is_properties_file | FALSE |
is_readme_file | FALSE |
is_language_file | TRUE |
has_private_key_extension | FALSE |
secret_length | 40 |
num_lines | 1 |
start_line | 212 |
start_column | 27 |
и подобными. Всего в нашей модели использовалось 30 характеристик.
Очевидной проблемой было отсутствие возможности передать бустингу суть контекста и место секрета в нём. Наиболее сложная задача при анализе кода – его токенизация. Токены, на которые разбивается текст кода, должны отражать его структуру, при этом игнорировать шум вроде пробелов и пропусков строки, а также понимать семантический смысл на разных языках программирования.
Чтобы добавить в нашу модель информацию об окружающем секрет контексте, мы попробовали пойти простым путём, выделили наиболее распространённые токены кода и добавили их TF-IDF признаки. Существенно улучшить качество при таком подходе не удалось. Поэтому мы стали изучать опыт решения подобных задач.
Среди небольших моделей рассмотрели также опыт применения свёрточных сетей для обнаружения личных данных авторизации в коде, описанный в статье. Авторы предлагают кодировать пароли посимвольно через One Hot представление, после чего обрабатывать их с помощью свёрточной архитектуры из 9 слоёв. В поиске паролей такой подход оправдан, поскольку пароль, придуманный человеком, отличается понятными схожими наборами символов, которые как раз выделяет свёрточная сеть. Для получения вектора контекста используется аналогичный посимвольный подход во избежание проблемы out-of-vocabulary.
Общая архитектура модели показалась нам интересной относительно больших языковых моделей, которые требуют куда более значительных вычислительных ресурсов. Однако для поиска секретов (в частности, API-токенов с более сложной структурой) такой подход не имел особого успеха по сравнению с бустингом. В связи с этим мы решили полностью обновить подход и стали смотреть в сторону трансформеров.
После изучения опыта применения трансформеров для такой задачи мы остановились на предобученной архитектуре CodeBERT. Эта архитектура требует относительно небольших ресурсов для инференса, при этом хорошо зарекомендовала себя в задаче выявления паролей в коде. Предобученная архитектура из статьи заточена только на пароли, при этом набор данных не раскрывается, из-за чего мы решили обучить свою модель.
Экспериментальным путём мы определили, что количество токенов контекста с секретом можно сократить до 256 токенов. Это объясняется тем, что потенциальный секрет часто указан при определении переменных, аргументов функций или атрибутов с характерными именами. Таким образом, для идентификации секрета достаточно рассмотреть несколько строчек кода. Сокращение числа токенов также позволило существенно сэкономить время и память при инференсе. Меньшего числа токенов оказывается недостаточно, поскольку сложноструктурированные секреты вроде API-токенов могут превращаться в очень длинную последовательность токенов.
Сбор данных для модели
Для обучения модели мы использовали набор данных SecretBench. Он содержит текст секрета, некоторые его характеристики, в том числе путь к скрипту и координаты самого секрета в нём. Из скриптов мы извлекли строки, содержащие секрет, вычистив случаи, когда из скрипта уже успели удалить секреты.
Наличие готового набора данных позволяет дообучить архитектуру CodeBERT для выявления секретов в коде, однако при детальном изучении в данных были обнаружены существенные недостатки. В SecretBench содержится довольно много дублей. Это одинаковые репозитории и скрипты с отличием лишь в коммитах. Нас интересовали только уникальные контексты с секретами. Кроме того, в данных встречались и случаи, когда при одном и том же контексте с секретом метка класса была то True Positive, то False Positive. Их пришлось вычищать вручную. После удаления всех дублей из 97.5 тысяч примеров осталось порядка 48 тысяч, некоторые из которых также выглядели довольно похожими (один и тот же секрет в разных фрагментах кода).
Чтобы увеличить количество уникальных примеров, мы использовали следующие аугментации:
замена существующего секрета, найденного по заданному паттерну, на новый, сгенерированный по аналогичному паттерну;
замена секретов категории Password на случайные пароли, взятые из публичных датасетов Password Strength Dataset и SecLists;
замена контекстов категории Authentication Key and Token на случайный контекст авторизации на том же языке программирования, что и исходный контекст. Примеры контекстов найдены в публичном датасете PassFinder.
Это позволило увеличить число примеров до 70 тысяч. Впоследствии часть из них была отсеяна на основе анализа ошибок модели.
Для оценки качества модели выделили около 24 тысяч случайных секретов таким образом, чтобы контрольная выборка не пересекалась с обучающей по репозиториям, а также не содержала аугментации. В нашем модуле секретов мы подсвечиваем пользователю не только найденные секреты, но и вероятность истинности его обнаружения. Поэтому нужно оценивать качество предсказанных моделью вероятностей. Кроме того, полученная выборка является несбалансированной: как правило, найденный кандидат не является секретом (около 85% случаев), то есть, преобладают ложноположительные срабатывания сканера.
Важно не пропустить ни одного реального секрета, при этом не увеличивать ручной труд на разбор пользователем всех срабатываний. В связи с перечисленными особенностями в качестве метрики оценки результатов мы выбрали площадь под Precision-Recall кривой (PR AUC). В нашем случае необходимо наиболее точно ранжировать результаты сканера, чтобы в первую очередь обращать внимание пользователя на реальные секреты. Такая метрика хороша и для несбалансированных данных, поскольку она отражает показатель качества при различных порогах округления вероятностей. Чем выше площадь под PR-кривой, тем точнее можно отделить ложноположительные срабатывания сканера от истинных секретов.
После обучения мы сравнили качество результатов бустинга и базового трансформера. Градиентный бустинг показывал около 0.70 PR AUC с дальнейшей просадкой качества в реальных условиях. По сравнению с этим решением базовый трансформер показал прирост PR AUC до 0.77.
Чтобы честно сравнить результаты классификации, мы округлили предсказанные вероятности по стандартному порогу 0.5 и заметили, что в результатах трансформера существенно сократилось число ложноотрицательных распознаваний моделью.
Подводные камни токенизатора
Анализ ошибок модели показал, что хуже всего определяются секреты, которые находятся в очень длинном контексте. Это приводило к тому, что в 256 токенах секрет мог вовсе не оказаться. Просмотрев, что возвращает предобученный токенизатор, мы также поняли, что большинство токенов – это сочетания по два символа. Это объясняется тем, что секреты выглядят как случайные сочетания символов, и это также увеличивает риск оставить секрет без внимания.
Из-за этой особенности мы переписали токенизатор таким образом, чтобы токены секрета оказались ровно посередине в случае, если его длина не превышает 256 токенов. Также сократили длинные секреты случайным удалением символов. Новая реализация включала три вызова токенизатора для контекста перед секретом, для самого секрета, а также для текста после него.
А привело всё это… к ухудшению качества до 0.72 PR AUC. Так мы выяснили, что контекст, указывающий на наличие секрета, в основном, находится перед текстом секрета. Переписав токенизатор так, чтобы он захватывал больше контекста “до”, нежели “после”, смогли улучшить качество до 0.86 PR AUC!
Узким горлышком оставались три вызова токенизатора вместо одного. Попробовали переписать это в один токенизатор, отслеживая позицию токенов секрета через опцию return_offsets_mapping=True.
Результат токенизации при этом стал отличаться от предыдущего варианта, поскольку теперь последние символы контекста до секрета могли “слиться” в один токен с первыми символами секрета. В варианте, где секрет проходил отдельную токенизацию, секрет обрабатывался моделью обособлено.
Рассмотрим пример строки:
‘SECRET 296866b2119ff5afbd84c4ee98dff791;/n’
При сплошной токенизации такого текста получим следующий набор ids токенов:
21536, 36995, 1132, 4671, 4280, 428, 2146, 1646, 3145, 245, 2001, 35470, 6232, 438, 306, 1942, 5208, 417, 3145, 5220, 134, 131, 50118
Теперь попробуем разбить исходную строку на текст до секрета, сам секрет и текст после секрета, проведём токенизацию каждого фрагмента в отдельности, после чего объединим ids обратно. Получим иной результат:
21536, 36995, 1437, 2890, 4671, 4280, 428, 2146, 1646, 3145, 245, 2001, 35470, 6232, 438, 306, 1942, 5208, 417, 3145, 5220, 134, 131, 50118
Здесь id 1437 соответствует пробелу, а 2890 – токену “29”. В то время как при токенизации исходного текста id токена 1132 соответствовал токену “ 29”.
Ради эксперимента оценили результат и при кодировании одним токенизатором. При несущественном различии в токенах получили изменение предсказаний с последующей просадкой качества по распознаванию истинно положительных примеров.
Чтобы выделить секрет и обрабатывать его токенизатором обособленно, добавили специальный токен <pad> до и после секрета и по нему выделили нужные 256 токенов.
'SECRET <pad>296866b2119ff5afbd84c4ee98dff791<pad>;/n'
На контрольной выборке вновь увидели небольшую просадку качества по сравнению с исходным подходом в три токенизации. Нашли расхождения: в некоторых примерах секрет мог встречаться несколько раз, и предыдущая версия токенизации первый раз обрабатывала текст секрета обособленно от контекста, а второй раз – вместе с самим контекстом. Это привело к мысли, что аугментацию текстов можно выполнять на уровне токенов: в нашем случае использовать такую “двойную” токенизацию для текста секрета.
Добавление двойной токенизации в длинных контекстах позволило увеличить количество распознаваний истинных секретов до 0.90 PR AUC!
При округлении по порогу 0.5 также видно существенное улучшение качества распознавания истинных секретов.
Полученное улучшение модели
Трансформеры позволили улучшить итоговое качество распознавания истинных секретов с 0.70 до 0.90 PR AUC. Этому предшествовала долгая работа с исходными данными, поиск новых примеров и аугментация. Наибольший интерес для дальнейшей работы представляет токенизация и выделение в качестве токенов общих шаблонов программирования, чтобы модель меньше зависела от выбранного пользователем языка и лучше понимала контекст. Как мы неоднократно замечали, прогноз модели может существенно измениться даже при изменении всего пары токенов в контексте. Таким образом, стоит обращать внимание на то, во что превращается текст после токенизации, действительно ли токены передают “осмысленную” информацию в модель или ей приходится работать с мусором в случае работы с непривычными для токенизатора текстами.
Надеемся, наш опыт был полезен и интересен. А как вы используете машинное обучение в безопасности? Будем рады обсудить в комментариях.
