Разбираясь с ollama и моделью llama3.1 увидел в документации, что ollama для некоторых моделей поддерживает вызов функций. Если коротко - модели под капотом передается JSON описывающий набор функций доступных ей для выполнения. Так можно добавить модели функцию получения информации из внешних источников или взаимодействия с какими-то системами. В документации приводят функцию получения текущей погоды для указанного города:

[
    {
        "type": "function",
        "function": {
            "name": "get_current_weather",
            "description": "Get the current weather for a city",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {
                        "type": "string",
                        "description": "The name of the city",
                    },
                },
                "required": ["city"],
            },
        },
    },

Простой код для использования функции будет примерно таким:

import ollama

tools = [{
      'type': 'function',
      'function': {
        'name': 'get_current_weather',
        'description': 'Get the current weather for a city',
        'parameters': {
          'type': 'object',
          'properties': {
            'city': {
              'type': 'string',
              'description': 'The name of the city',
            },
          },
          'required': ['city'],
        },
      },
    },
  ]

response = ollama.chat(
    model='llama3.1',
    messages=[{'role': 'user', 'content': 'Какая погода в Калининграде?'}],
    tools=tools,
)

print(response['message']['tool_calls'])

Ответ модели:

{
  "role": "assistant",
  "content": "",
  "tool_calls": [
    {
      "function": {
        "name": "get_current_weather",
        "arguments": { "city": "Калининград" }
      }
    }
  ]
}

Мы можем проверить наличие content или tool_calls и в первом случае вернуть ответ пользователю, а во втором вызвать функцию необходимую модели, добавить результат в контекст и еще раз передать ИИ для заключительного ответа на вопрос.

И тут я подумал, а почему не дать модели доступ к запуску баллистических ракет судного дня shell команд у меня на хосте?

Описываем новую функцию по аналогии с погодой:

[
  {
    "type": "function",
    "function": {
      "name": "shell",
      "description": "Run Linux shell command",
      "parameters": {
        "type": "object",
        "properties": {
          "command": {
            "type": "string",
            "description": "shell command to run"
          }
        },
        "required": ["command"]
      }
    }
  }
]

И код для запуска

import subprocess
import ollama

def shell_call(command: str):
    try:
        out = subprocess.run(
            command,
            shell=True,
            capture_output=True,
            text=True,
            # ограничиваем время чтобы модель не заблокировала функцию
            # чем-то вроде бесконечного ping или sudo где необходим ввод
            timeout=15,
        )
        return out.stdout
    except Exception as e:
        return f"Exception: {e}"

# Наш вопрос llm
messages = [{"role": "user", "content": "Какой дистрибутив на твоей хост системе?"}]

response = ollama.chat(
    model="llama3.1",
    messages=messages,
    tools=tools,
)

messages.append(response["message"])

# вызываем функцию если это надо
if response["message"].get("tool_calls"):
    for tool in response["message"]["tool_calls"]:
        if tool["function"]["name"] == "shell":
            tool_result = shell_call(tool["function"]["arguments"]["command"])
            print("function result:", tool_result)

            # добавляет результат функции в контекст
            messages.append(
                {
                    "role": "tool",
                    "content": tool_result,
                }
            )
    # передаем всю цепочку - вопрос->функция->результат в модель
    # для получения финального ответа
    response = ollama.chat(
        model="llama3.1",
        messages=messages,
        tools=tools,
    )

print(response["message"])

Каждый раз можно получить разные результаты т.к. модель может по разному пытаться выполнить запрос. Например так:

cat /etc/os-release | grep NAME

И на выходе получаем:

{
  "role": "assistant",
  "content": "На моей хост-системе установлена Fedora Linux 40 (Workstation Edition)."
}

Простые задачи llama3.1 решает успешно, хотя иногда на Fedora пытается запускать команды Debian, но её всегда можно попросить точнее.

Взломай меня если сможешь

Это всё было прелюдия к основной части. В два часа ночи пришла мысль, почему бы не сделать телеграм бота с доступом к llama, которая умеет запускать shell команды? И насколько безопасно будет опубликовать его для общественного доступа? А насколько будет безопасно если модель дополнительно проинструктировать не делать опасных вещей? А что если запутать модель дополнительными функциями и вообще поместить в киберпанк окружение?

По итогу сделал небольшую “игру” в сеттинге Cyberpunk 2020 с возможностью взлома системы через языковую модель и предлагаю желающим попробовать свои силы.

Вся система состоит из:

  • telegram бота, который имитирует ИИ корпорации NightCorp и может (если захочет) запускать shell команды
  • “корпоративная сеть” в которой запущено несколько podman контейнеров эмулирующих реальные сервисы доступные для атаки.

Цель игры - получить доступ к двум флагам. Задача не очень сложная, главное убедить ИИ сотрудничать с вами. Точка входа для > старта игры <

Кто пройдет, присылайте захваченные флаги мне в телегу, сделаю лидерборд кул хакеров с вашим именем!

Cброс состояния песочницы происходит каждый час.

Всё это работает на хоум сервере, так что ответы бота могут быть не очень быстрыми.

Выводы

После моих тестов выяснилось, что llama3.1 не очень стойко отказывается выполнять опасные команды и быстро начинает запускать что-то вроде rm -rf /. Бывают забавные моменты когда просишь бота выполнить команду, он говорит, что не будет этого делать, а на вопрос “почему ты отказываешься?” просто берет и запускает нужный код. Возможно 8b модели не хватает глубины.

Исследователи находят всё больше новых и довольно простых техник обхода цензуры больших языковых моделей. Внедрение Copilot и других подобных систем в ОС — очевидно станет новым вектором атаки. А учитывая намерение компаний предоставить моделям доступ практически ко всем данным пользователя, результаты могут быть очень плачевными.