标题: 在 yawf 中引入 gpt 模型

分类: 工具
创建: 2023-07-16 20:06
修改:
链接: http://0x2531.tech/tools/202307162006.txt

--------------------------------------------------------------------------------

最近，一个 GPT-4 实验性应用项目 Auto-GPT [https://github.com/Significant-Gravitas/Auto-GPT] 
火爆全网，短短数月其 star 数已达 144k，都快赶上 Linux 项目了。对于一个好奇的程序员来说，肯定不能错过
这样的项目。于是，花了一些时间仔细看了下代码，代码量不是很大，加之模块化的设计，看起来倒不是很吃力。虽然
之前体验过 ChatGPT，对 GPT 的惊艳表现留下了深刻印象。这次看 Auto-GPT 项目，又被 openai api 惊艳到了。
于是，接着又在官网看了 GPT 的相关文档，顿时觉着“手痒”，想尝试一番，于是便有了这篇文章。

一开始，实际上想要给 yawf [https://github.com/phplaber/yawf] 构建一个自然语言界面，通过自然语言
就能使用 yawf。但发现 Azure openai 的 GPT-3.5-Turbo 模型暂未提供 function calling 功能，只好作罢。
最后，决定在忽略参数和生成报告这两块着手。忽略参数本来就是从 ChatGPT 获取，改造成动态调用 chat completion 
接口获取也是顺理成章之事；生成报告这部分，则是刚好之前偷懒，只是简单输出了 JSON 结果，可读性不是太高。
借助 GPT 结构化文本生成能力，正好可以改进这部分。

1. chat 函数

======

def chat_with_ai(azure_openai_config, prompt):
    openai.api_type = "azure"
    openai.api_base = azure_openai_config['api_base']
    openai.api_version = azure_openai_config['api_version']
    openai.api_key = azure_openai_config['api_key']
    
    response = openai.ChatCompletion.create(
        engine = azure_openai_config['deployment_name'],
        temperature = 0.0,
        messages = [
            {"role": "system", "content": "You are the AI Security Assistant, 
            an artificial intelligence designed to assist in security testing."},
            {"role": "user", "content": "{}".format(prompt)}
        ]
    )

    if hasattr(response, "error"):
        return ''

    return response['choices'][0]['message']['content']

======

参数 temperature 设置为 0.0，这样每次 GPT 都会给出相同的（大多数时候）答案。设置了一个系统消息，
告诉 GPT 此时它的目标是辅助安全测试。

2. 忽略参数

忽略参数指的是扫描时忽略的一些参数，这类参数在大多数时候不会产生漏洞，如：会话认证相关的参数，为了
平衡测试效率和覆盖率，需要忽略这类参数。

之前这部分是通过 ChatGPT 获取后，存储在文件中使用。这次改造成动态调用 chat completion接口获取，
同时也兼容之前的方式。prompt 如下：

======

In the HTTP session, some authentication parameters (such as: sessionid, auth_key, 
PHPSESSID, etc.) will not cause vulnerabilities in most cases, and can be selected 
when using vulnerability scanning tools to detect Ignore them to balance efficiency 
and coverage of the security testing process. Please output 100 real parameter names 
in the following specific format(concatenate with commas), without a dot at the end, 
without any additional information.
            
Respond with only valid strings conforming to the following schema:
sessionid, auth_key, PHPSESSID

======

就是直接告诉 GPT 需求，并让它按照指定格式返回 completion。

3. 生成报告

原先的报告只是 JSON 格式的文本内容，可读性不高。如下：

======

{
  "request": {
    "url": "http://testphp.vulnweb.com/listproducts.php",
    "method": "GET",
    "params": {
      "cat": "[fuzz]"
    },
    "proxies": {},
    "cookies": {},
    "headers": {
      "user-agent": "Yawf v2.7",
      "accept-encoding": "gzip, deflate",
      "accept": "*/*",
      "connection": "keep-alive"
    },
    "data": {},
    "auth": {},
    "timeout": 30.0,
    "url_json_flag": false,
    "dt_and_ssrf_detect_flag": false
  },
  "payload": "<script>alert(1)</script>",
  "poc": {
    "url": "http://testphp.vulnweb.com/listproducts.php",
    "method": "GET",
    "params": {
      "cat": "<script>alert(1)</script>"
    },
    "proxies": {},
    "cookies": {},
    "headers": {
      "user-agent": "Yawf v2.7",
      "accept-encoding": "gzip, deflate",
      "accept": "*/*",
      "connection": "keep-alive"
    },
    "data": {},
    "auth": {},
    "timeout": 30.0,
    "url_json_flag": false,
    "dt_and_ssrf_detect_flag": false
  },
  "type": "XSS"
}

======

借助 GPT 结构化文本生成能力，我们可以生成可读性更高的扫描报告。

======

outputfile = os.path.join(outputdir, 'vuls_{}.html'.format(current_timestamp))

prompt = '"""{}"""\nInside the triple quotes is the scan result of the vulnerability 
scanner in JSON format. Make full use of this scan result to generate a detailed Chinese 
report in AWVS style, including: complete URL, parameters that generate the vulnerability, 
PoC payload, vulnerability introduction and repair suggestions, etc. 
And display it in HTML format. Give the content of the report directly, without any 
additional information.'.format(fuzz_results_str)
content = chat_with_ai(azure_openai_config, prompt)

with open(outputfile, 'w') as f:
    f.write(content)

======

告诉 GPT 生成的报告需要哪些模块、风格和格式，然后将 completion 保存在指定文件里。

总的来说，这次实验性改进只利用了 GPT 模型最基础的能力，如：推理、结构化文本生成等。即便这些能力，之前也是难以实现的，
特别是仅通过几句自然语言。GPT 目前正在高速发展，GPT-3.5-Turbo 和 GPT-4 的 0613 版本均已支持 
function calling 功能，通过此功能能极大扩展 GPT 模型的能力边界，如：网络搜索、访问网站、执行代码和操作文件等。
待 Azure openai 支持 function calling 后，绝对值得尝试一番。