写在前面
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 BaseToolimport subprocessclass 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 = Truedef _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.returncodeasync 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, AgentTypefrom xray_tools import XrayScanToolllm = ChatOpenAI(model_name="gpt-3.5-turbo-0613",temperature=0,# 填入自己的ChatGPT APIopenai_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的一次实践探索，还是很成功的，只能说魔法加在子弹上，不可估量。
