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 Anyfrom 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 ChatOpenAIfrom 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()]# 初始化agentagent = 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的一次实践探索，还是很成功的，只能说魔法加在子弹上，不可估量。