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的代码

python 复制代码
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对工具进行调用

python 复制代码
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
)

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

python 复制代码
resp = agent.run("扫描一下http://127.0.0.1/,看看有没有敏感文件目录")

print(resp)

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

写在后面

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

相关推荐
deephub8 小时前
AI代理性能提升实战:LangChain+LangGraph内存管理与上下文优化完整指南
人工智能·深度学习·神经网络·langchain·大语言模型·rag
都叫我大帅哥13 小时前
LangChain分层记忆解决方案:完整案例
python·langchain
alex10013 小时前
AI Agent开发学习系列 - langchain之LCEL(5):如何创建一个Agent?
人工智能·python·语言模型·langchain·prompt·向量数据库·ai agent
liliangcsdn16 小时前
mac测试ollama llamaindex
数据仓库·人工智能·prompt·llama
青Cheng序员石头17 小时前
Prompt Engineering vs Vibe Coding vs Context Engineering
langchain·llm·aigc
数据智能老司机17 小时前
构建由 LLM 驱动的 Neo4j 应用程序——使用 Neo4j 和 Haystack 实现强大搜索功能
langchain·llm·aigc
都叫我大帅哥1 天前
🚀 LangGraph终极指南:从入门到生产级AI工作流编排
python·langchain
showyoui2 天前
LangChain vs LangGraph:从困惑到清晰的认知之路(扫盲篇)
langchain·ai编程
_一条咸鱼_2 天前
LangChain记忆序列化与持久化方案源码级分析(37)
人工智能·面试·langchain
数据智能老司机2 天前
构建由 LLM 驱动的 Neo4j 应用程序——揭开 RAG 的神秘面纱
langchain·llm·aigc