AI驱动的安全扫描工具：HackBot

原创

极致攻防实验室

发布于 2023-07-01 00:05:00

2820

发布于 2023-07-01 00:05:00

文章被收录于专栏：极致攻防实验室极致攻防实验室

免责声明：

本文章或工具仅供安全研究使用，请勿利用文章内的相关技术从事非法测试，由于传播、利用此文所提供的信息而造成的任何直接或者间接的后果及损失，均由使用者本人负责，极致攻防实验室及文章作者不为此承担任何责任。

1.前言

OpenAI在6月13日在API级别发布了新的更新,除了发布新的模型gpt-3.5-turbo-16k和更新了模型gpt-3.5-turbo为gpt-3.5-turbo-0613，最有意思的还是gpt-3.5-turbo-0613模型不仅input费用降低了25%，还支持了一个新的功能：Fucntion calling。

本文将使用Golang语言，利用OpenAI的Function calling功能与Projectdiscovery武器库进行结合，制作一个用于安全扫描的工具：HackBot 。

2. Fucntion calling（函数调用）

2.1 什么是Fucntion calling

Function calling是可以让我们用自己的函数当作调用chatgpt的参数，在函数中我们可以做任何事情，例如获取网络上的数据，查询自己的数据库等。

2.2 为什么用Fucntion calling

使用过chatgpt的同学都知道，chatgpt没有实时的网络数据，特别是在安全领域某些问题会存在安全限制，而Fucntion calling正好可以解决这一问题，有了Function calling我们可以在函数中调用最新的数据返回给chatgpt，甚至让chatgpt去判断接下来应该调用什么函数。

2.3 官方简单的例子

官网的天气查询例子:How_to_call_functions_with_chat_models

3. 使用Fucntion calling进行安全扫描

下面将使用Go语言，利用OpenAI的Function calling功能与projectdiscovery武器库进行结合，制作一个用于安全扫描的概念小工具：HackBot 。

3.1 在Golang中调用ChatGPT

使用github.com/sashabaranov/go-openai可以很方便去调用chatgpt

package main

import (
  "context"
  "fmt"
  openai "github.com/sashabaranov/go-openai"
)

func main() {
  client := openai.NewClient("your token")
  resp, err := client.CreateChatCompletion(
    context.Background(),
    openai.ChatCompletionRequest{
      Model: openai.GPT3Dot5Turbo,
      Messages: []openai.ChatCompletionMessage{
        {
          Role:    openai.ChatMessageRoleUser,
          Content: "Hello!",
        },
      },
    },
  )

  if err != nil {
    fmt.Printf("ChatCompletion error: %v\n", err)
    return
  }

  fmt.Println(resp.Choices[0].Message.Content)
}

针对于国内需要代理才能访问chatgpt，我们可以自定义一个代理的httpClient

func (g *GptClient) InitClient(opt *Options) {
  config := openai.DefaultConfig(opt.AutoToken)

  if opt.Proxy != "" {
    proxyURL, err := url.Parse(opt.Proxy)
    if err != nil {
      fmt.Println("Error:", err)
      os.Exit(1)
    }

    transport := &http.Transport{
      Proxy: http.ProxyURL(proxyURL),
    }

    httpClient := &http.Client{
      Transport: transport,
    }
    config.HTTPClient = httpClient
  }

  g.Client = openai.NewClientWithConfig(config)
}

3.2 在Golang中使用ChatGPT函数调用功能

首先我们定义一个利用Projectdiscovery httpx进行网站存活探测与信息收集的函数

httpx.go:


package core

import (
  "fmt"
  "github.com/projectdiscovery/goflags"
  "github.com/projectdiscovery/gologger"
  "github.com/projectdiscovery/gologger/levels"
  "github.com/projectdiscovery/httpx/runner"
)

type Httpx struct {
  Url string `json:"url"`
}

func (h *Httpx) Run() (content string, err error) {
  gologger.DefaultLogger.SetMaxLevel(levels.LevelSilent)
  options := runner.Options{
    Methods:         "GET",
    InputTargetHost: goflags.StringSlice{h.Url},
    FollowRedirects: true,
    Silent:          true,
    OnResult: func(r runner.Result) {
      // handle error
      if r.Err != nil {
        err = r.Err
        return
      }
      content = fmt.Sprintf("Url:[%s] Host:[%s] Title:[%s] Status:[%d] Length:[%d] Favicon:[%s] Technologies:%v",
        r.URL, r.Host, r.Title, r.StatusCode, r.ContentLength, r.FavIconMMH3, r.Technologies)
    },
  }

  if err = options.ValidateOptions(); err != nil {
    return
  }

  r, err := runner.New(&options)
  if err != nil {
    return
  }
  defer r.Close()

  r.RunEnumeration()

  return
}

然后再根据go-openai提供好的FunctionDefinition格式，去编写一个“HttpxFunc”

func.go:

var HTTPXFunc = openai.FunctionDefinition{
  Name:        "httpx",
  Description: "Used for website survival detection and basic website information acquisition, such as website title, HTTP response content, etc.(用于网站存活探测和网站基本信息获取，如网站标题、http响应内容等)",
  Parameters: &openai.JSONSchemaDefinition{
    Type: openai.JSONSchemaTypeObject,
    Properties: map[string]openai.JSONSchemaDefinition{
      "url": {
        Type:        openai.JSONSchemaTypeString,
        Description: "destination address(url)",
      },
    },
    Required: []string{"url"},
  },
}

var MyFunctions = []openai.FunctionDefinition{
  HTTPXFunc,
}

根据 go-openai修改附带调用functions的发送请求的代码：

resp, err := g.Client.CreateChatCompletion(
  context.Background(),
  openai.ChatCompletionRequest{
    Model:        openai.GPT3Dot5Turbo0613,
    Messages:     g.MessageHistory,
    Functions:    MyFunctions,
    FunctionCall: "auto",
  },
)

if err != nil {
  return
}

respMessage := resp.Choices[0].Message
if respMessage.FunctionCall == nil {
  fmt.Println("Error: The input information is too little to call the function")
  return
}

funcName := resp.Choices[0].Message.FunctionCall

content, err := FuncCall(funcName)
if err != nil {
  return
}

当chatgpt接受到我们的问题，将会返回需要调用的函数名，所以需要根据函数名去执行相应的函数。

func.go:

func FuncCall(call *openai.FunctionCall) (content string, err error) {
  switch call.Name {
  case "httpx":
    var f Httpx
    if err = json.Unmarshal([]byte(call.Arguments), &f); err != nil {
      return
    }
    content, err = f.Run()
  }
  return
}

效果(Httpx)：