写在前面
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的一次实践探索，还是很成功的，只能说魔法加在子弹上，不可估量。
