Привет, хабр! В этой статье я расскажу о новом, весьма интересном способе заражения пользователя с помощью Tensorflow (кстати, не только TF подпадает под это - но и все библиотеки для работы с ML)

Часть 0. Предисловие.

Я делаю это публичным, ведь гугл закрыли мой отчет с вот такими словами:

Hi,

Thanks for the report, but we do not consider this a valid security bug.
Creating a typo-squatted malicious library has nothing to do with TensorFlow. The scenario you described can be abused against any other package.

Thanks again for your report and time,
The Google Bug Hunter Team

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

И кстати: использование любой уязвимости кроме как для обучения незаконно.

Часть 1. Введение в суть атаки.

Сама суть атаки проста - что если заразить не пользовательские модели, а саму среду пользователя? То есть статические и динамические библиотеки (.pyd, .so, .dll). Это не так сложно, как кажется, но это дает весьма большие возможности. Ведь человек скачавший зараженный пакет не заметит никакой разницы с официальным - он будет иметь одинаковое название, версию, только вот хэш будет различаться. (кстати, когда последний раз сверяли хэши файлов с официальными?). Главная проблема почему это до сих пор работает - отсутствие подписи бинарных файлов. Хотя, существует простая, на первый взгляд, защита - RECORD файл, но проблема в том, что его не так сложно модифицировать. Подробнее - будет дальше.

Часть 2. Объяснение вектора атаки.

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

2.1 Использование уязвимости через атаку на цепочку поставок

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

Шаг 1. Создание:

Берется официальный пакет, например с dev зеркала (качали когда-нибудь nightly torch?) последней версии и исходный код с github. Обычно в .ci репозиториев есть полноценные скрипты для сборки. Поэтому, после добавления полезной нагрузки в место, которое может вызываться при любом взаимодействии юзера или при определенных условиях, собрать проект не составит труда. Потом злоумышленник берет созданный .whl и достает из него лишь необходимые библиотеки с компилированным кодом - .so, .dll, .pyd. И как итог - получается пакет с виду одинаковый с оригиналом. Но различающийся буквально в паре файлов. Чуть подробнее про редактирование RECORD файла: необходим лишь простой скрипт на python, который исходя из названий найдет файлы в папке и в самом файле. И просто заново пересчитает sha256 хэш для них.

Шаг 2. Распространение:

Все не так сложно как кажется. Не надо делать на pypi пакеты с похожими названиями, надеясь, что пользователь по ошибке установит вредоносный пакет. Все намного проще - использование социальной инженерии. Злоумышленник ищет потенциальных жертв на StackOverflow или в Issues на Github и предлагает, например, такое заманчивое предложение для решения возникшей проблемы:

У меня на этой новой dev версии TF появился прирост скорости на 5090 и я проверил твою ошибку - у меня она не возникает. Вот, держи ссылку:

Либо можно дать сборку с "заряженным" кодом (Веская причина не доверять сборкам из непроверенных источников, не так ли?)

Шаг 3. Исполнение атаки:

Пользователь либо ставит зараженный пакет, либо запускает портативную сборку - и автоматически происходит заражение, ведь с++ часть кода уже загружена в память и выполнена. Что насчет проверки на стороне pip? Она полностью пройдена, ведь RECORD файл валидный, а подписей файлов нет, следовательно ошибки не возникнет.

Часть 3. Реальный пример работы.

Я не буду изощряться в методе доказательства работы - обойдусь весьма простой и рабочей штукой, null ptr (но реальный злоумышленник может использовать что угодно - потому это RCE):

*(int*)0 = 0;

Патчить я буду tensorflow/core/platform/cpu_feature_guard.cc - это идеальное место для встраивание полезной нагрузки. Загружается всегда.

Я встроил код в данном месте, буквально перед #ifdef SSE . После этого процесс утомительной сборки длительностью пару часов при помощи официальных скриптов в папке .ci. Я получаю файл который итак весьма похож на официальный из dev канала, но я для своего удобство сделал замену dist-info, как и говорил чуть раньше. Итак, вот работа TF с официального DEV whl:

(venv) whoami@root:/home/whoami/tf$ python -c "import tensorflow"
I0000 00:00:1756730745.507629   67483 cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: SSE3 SSE4.1 SSE4.2 AVX AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.

А вот работа с зараженной версии под видом этого же колеса:

(venv) whoami@root:/home/whoami/tf$ python -c "import tensorflow"
Trace/breakpoint trap (core dumped)

Мне кажется существенная разница. Не так ли?

Часть 4. Как защититься от этого?

Пока не реализована подпись бинарных файлов и проверка их с помощью pip - придется делать все своими силами. Но рекомендации сами по себе достаточно просты:

1. Не лениться сравнивать хэш. "Знакомый" поделился с тобой библиотекой, которую собрал сам? Или добрый человек дал тебе решение твоей проблемы, показав на своей версии библиотеки? Сравни хэши с официальными, если найдешь пакет один в один с таким названием. Это не так сложно, как кажется!

2. Не доверяй портативным сборкам. Как бы это не было удобно, но лучше попроси скрипт установки, а не используй портативную сборку. Простое действие - но отсекает множество возможных проблем (и да, файл установки тоже проверяй! вдруг там тоже запрячется странный пакет, который может быть заражен?)

Часть 5. Чуть-чуть больше насчет подписей кода и pip.
На самом деле интересно, как так вышло, что комьюнити питона не придумали защиту от такого? Хотя, казалось бы, подпись кода решила бы эту проблему навсегда. Но если этого нет - значит проблема глубже.

Исторически pypi имел возможность добавлять GPG-подписи, но это был ручной процесс, да и pip это проверять не умеет. Сегодня сообщество движется в сторону внедрения TUF (The Update Framework). Подробнее он описан в PEP 458, он как раз и предназначен для предотвращения такого рода атак, но увы, на момент написания данной статьи интеграция его в pip еще на завершена и не работает по умолчанию.

Поэтому, до сих пор можно заразить конечного пользователя лишь просто подменив RECORD файл, что весьма печально. Но есть встроенная защита - pip install --hash , но будем честны, вы хоть раз пользовались этим? Лично я - нет. Да и в добавку это требует ручных действий от пользователя, а не работает автоматически.

Часть 6. Подводя итоги.

На самом деле сама идея атаки не так нова - просто ее никогда не поднимали. И все вслепую доверяли пакетам, которые ставят. Хотя, это не лучшая практика. Будьте в безопасности и не дайте себя взломать. А если предупрежден - значит вооружен.

P.S. Хотел бы услышать и ваше мнение насчет данной атаки. И да, я люблю отвечать на вопросы =)