langchain+xray：prompt控制漏洞扫描

写在前面

xray是长亭推出的一款漏洞扫描工具。

langchain是调用LLM大模型完成自动化任务的框架。

本篇文章是对langchain自定义工具的探索，通过编写一个xray调用的工具，联合ChatGPT对xray进行调用，实现对目标的漏洞扫描。

xray功能分析

首先分析一下xray的功能

爬虫扫描

xray webscan --basic-crawler http://example.com --html-output vuln.html

主动扫描

xray webscan --url http://example.com/?a=b --html-output single-url.html

指定插件扫描

xray webscan --plugins cmd-injection,sqldet --url http://example.com

分析之后，因为我想用一个工具实现，而不用每个功能都写一个工具。于是，设想了三个参数

target：扫描目标
vul_type：扫描漏洞类型
mode：扫描方式

工具编写

直接上tools的代码

from typing import Any
from langchain.tools import BaseTool

import subprocess


class XrayScanTool(BaseTool):
    name = "xray_vul_scan"
    description = "使用Xray漏洞扫描器，对目标（参数 target）进行漏洞扫描，可以指定全漏洞扫描，也可以针对性指定（参数vul_type）。" \
    "漏洞类型包含：xss, sqldet, cmd-injection, dirscan, path-traversal, xxe, phantasm, upload, brute-force," \
    "jsonp, ssrf, baseline, redirect, crlf-injection, xstream。通过list传入" \
    "还可以指定扫描方式（参数mode），爬虫扫描crawler，还是非爬虫扫描scan"

    # 指定xray的路径
    XRAY_PATH = "./xray_windows_amd64.exe"
    return_direct = True

    def _run(self, target, vul_type=["all"], mode="scan") -> Any:
        if mode == "crawler":
            base_cmd = f"{self.XRAY_PATH} webscan --basic-crawler {target} "
        else:
            base_cmd = f"{self.XRAY_PATH} webscan --url {target} "

        if "all" not in vul_type:
            base_cmd += "--plugins "
            for vul in vul_type:
                base_cmd += vul + ","
            base_cmd = base_cmd.rstrip(",")

        base_cmd += " --html-output vuln.html"
        
        if os.path.exists("vuln.html"):
            os.remove("vuln.html")
            
        result = subprocess.run(base_cmd.split(" "))
        return result.returncode

    async def _arun(self, *args: Any, **kwargs: Any,) -> Any:
    	raise NotImplementedError("This tool does not support async")

可以看到，我在描述中说的很详细，主要是为了能让ChatGPT可以准确的传入我需要的参数，只要描述足够详细，剩下的你可以完全相信ChatGPT！

然后，下面开始创建agent对工具进行调用

from langchain.chat_models import ChatOpenAI
from langchain.agents import initialize_agent, AgentType

from xray_tools import XrayScanTool

llm = ChatOpenAI(
    model_name="gpt-3.5-turbo-0613",
    temperature=0,
    # 填入自己的ChatGPT API
    openai_api_key=API_KEY,
    openai_api_base=BASE_ADDRESS
)

# 构建工具列表
tools = [XrayScanTool()]

# 初始化agent
agent = initialize_agent(
    agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
    tools=tools,
    llm=llm,
    verbose=True
)

这部分不多赘述，下面进行测试

resp = agent.run("扫描一下http://127.0.0.1/，看看有没有敏感文件目录")

print(resp)

可以看到，完美传入了我们所需的参数，并进行了执行。

写在后面

这个demo只是一次对langchain tools的一次实践探索，还是很成功的，只能说魔法加在子弹上，不可估量。