基于ChatGPT打造安全脚本工具流程

前言


以前想要打造一款自己的工具，想法挺好现实上是难以实现，第一不懂代码的构造，只有一些工具脚本构造思绪，第二总是像重复造轮子这种繁琐枯燥工作，抄抄改改搞不清晰逻辑，想打造一款符合自己工作的自定义的脚本工具难度倍增，但是随着AI的鼓起极大的改变了我们的生活工作方式，只需利用ChatGPT即可打造自己的专属自定义工具，下面作者以自己的几个脚本工具的案例为例，报告一下如何利用ChatGPT开发安全脚本工具。（注：作者专注安全其他方向，非正儿八经的安全开发，纯粹解说思绪，工具大抵之际，纯粹不入流勿喷）

案例一 漏洞简报生成


工具前景


作者是一名渗透测试工程师，在工作中避免不了写漏洞报告，为了提高工作效率，避免繁琐的文档工作，所以有一下脚本工具诞生，工具潦草但实用，快速生成一份漏洞简报，只需填写对应的漏洞结果，避免了上网查询漏洞相关细节和修复方式，基于web.xlsx表格内容为存储漏洞信息，有一定局限性，但工具目前可塑性很强，有待提高。

工具名称


本工具为《Bug_writer》漏洞报告辅助工具，该工具旨在提高渗透测试职员，和网络安全从业者，爱好者等人群的工作效率,使得更快速准确地填写更多的漏洞报告，虽然作者不是专业开发,但作为一名安全从业者，作者在专心制作这个工具来解决自己的工作需求，之前Github有更新本项目，后来参加HVV大意了没有闪被溯源了，Github废弃了，马甲相关信息均废弃（还不是因为丢人）。

工具结构

1. 工具结构
2. Bug_writer:
3. │  
4. │  main.py
5. │  README.md
6. │  requirements.txt
7. │
8. ├─config
9. │  │  WEB.docx
10. │  │  WEB.xlsx
11. │  │
12. │  └─漏洞模版
13. │          简报模版一.docx
14. │          默认模板.docx
15. │
16. └─output
18. 说明：
20. 主目录 mian.py为脚本文件
21. │  README.md 为使用说明
22. │  requirements.txt 为需要安装的第三方库
23. │  config文件夹下
24. │    WEB.docx 简报模版
25. │    WEB.xlsx 漏洞库
26. │    漏洞模版文件夹下，为模版样式
27. │  output文件夹为输出的简报文件

复制代码

工具思绪


在此之前我们是参考xlsx漏洞模板，用于填写漏洞报告，后来转念一想，可以把漏洞里面的信息直接填充就好了，随机一发不可收拾的开始搞这个脚本工具，现将模版修改如下所示，然后开始构思流程

首先我的工具定位是可以生成简报的脚本，其次我对它要求为代码简朴快速生成且可塑性强，然后如下所示

思量python语言较为方便，使用python作为脚本工具开发语言，优点在于它的第三方库和语言通用性

第一步


1.先将xlsx模版放在config文件夹中，首先需要生成pyhton脚本模糊查询漏洞名称，假设为sql注入，我们输入注入，该工具会在xlsx表格中查询漏洞名称列，如果没有查询到存在注入二字结果，直接结束返回下令行页面告知未查询结果，如果有该漏洞名存在，举行下一步，根据上述的思绪，编写文案，在Chatgpt上执行，会得出初代源码，根据初代源码举行修改不足，完成上述所示需求。

第二步


2.当初代源码可以调用模糊查询时，我们就需要将其他列的信息按照格式展示在下令行，如何展示可以在ChatGpt上美满我们第二步需求,如下所示，这里已经美满了上述需求.

代码如下

1. import argparse
2. import os
3. import pandas as pd
6. class VulReportGenerator:
7.     def __init__(self):
8.         self.parser = argparse.ArgumentParser(description='漏洞报告生成器')
9.         self.parser.add_argument('-c', type=str, help='指定模糊查询漏洞名称.')
10.         self.parser.add_argument('-t', action='store_true', help='生成报告')
11.         self.args = self.parser.parse_args()
12.         self.df = pd.read_excel('./config/WEB.xlsx')
14.     def search_vuls(self):
15.         search_name = self.args.c.lower()
16.         result = self.df[self.df['漏洞名称'].str.lower().str.contains(search_name)]
17.         result_count = result.shape[0]
19.         if result.empty:
20.             print('<<这里并没有发现结果--您可以去查看是否输入错误>>')
21.         else:
22.             print(f'查询到 {result_count} 条记录：')
24.             for index, row in result.iterrows():
25.                 print('漏洞名称:  {}\n风险级别: {}\n漏洞描述:\n\t {} \n加固建议: \n\t {}\n\n\n'.format(
26.                     row['漏洞名称'], row['风险级别'], row['漏洞描述'], row['加固建议']))
28.             if self.args.t:
29.                 self.generate_report(result)
31.     def generate_report(self, data):
32.         # 在这里编写生成报告的逻辑，例如将结果输出到文件或者其他操作
33.         print("生成报告：" + str(data))
36. if __name__ == '__main__':
37.     generator = VulReportGenerator()
38.     generator.search_vuls()

复制代码

第三步


这一步的时间，我们可以看到，样式有些丢脸，我们可以使用图形化字符工具生成一个字符图形，然后再美满增加颜色样式。

代码如下

1. import argparse
2. import os
3. import pandas as pd
5. # 创建命令行参数解析器
6. print('''
7. \033[31m  ____                              _ _            \033[0m
8. \033[31m | __ ) _   _  __ _  __      ___ __(_| |_ ___ _ __ \033[0m
9. \033[31m |  _ \| | | |/ _` | \ \ /\ / | '__| | __/ _ | '__|\033[0m     
10. \033[31m | |_) | |_| | (_| |  \ V  V /| |  | | ||  __| |   \033[0m     
11. \033[31m |____/ \__,_|\__, |   \_/\_/ |_|  |_|\__\___|_|   \033[0m     Version: \033[33m公测版本 \033[0m
12. \033[31m              |___/\033[0m                                     Tools: \033[33m漏洞报告辅助工具\033[0m      
13.             \033[34mgithub.com/yichensec/Bug_writer\033[0m     
15. 一款面向用于安服，渗透测试人员，网络安全从业者等人群的漏洞报告辅助工具可以生成漏洞测试简报，可自己私人定制。''')
16. print('''\033[33m================================== 渗透漏洞报告书写员的开始 =========================================\033[0m''')
18. class VulReportGenerator:
19.     def __init__(self):
20.         self.parser = argparse.ArgumentParser(description='漏洞报告生成器')
21.         self.parser.add_argument('-c', type=str, help='指定模糊查询漏洞名称.')
22.         self.parser.add_argument('-t', action='store_true', help='生成报告')
23.         self.args = self.parser.parse_args()
24.         self.df = pd.read_excel('./config/WEB.xlsx')
26.     def search_vuls(self):
27.         search_name = self.args.c.lower()
28.         result = self.df[self.df['漏洞名称'].str.lower().str.contains(search_name)]
29.         result_count = result.shape[0]
31.         if result.empty:
32.             print('<<这里并没有发现结果--您可以去查看是否输入错误>>')
33.         else:
34.             print(f'查询到 {result_count} 条记录：')
36.             for index, row in result.iterrows():
37.                 print('\033[0m\033[33m漏洞名称:\033[0m\033[32m  {}\n\033[0m\033[33m风险级别:\033[0m\033[31m  {}\033[0m\n\033[33m漏洞描述:\033[0m  \n\t\033[34m  {}\033[0m\n\033[33m加固建议:\033[0m\033[35m \n\t {}\033[0m\n\n\033[33m=====================================================================================\033[0m\n'.format(row.漏洞名称, row.风险级别, row.漏洞描述, row.加固建议))
39.             if self.args.t:
40.                 self.generate_report(result)
42.     def generate_report(self, data):
43.         # 在这里编写生成报告的逻辑，例如将结果输出到文件或者其他操作
44.         print("生成报告：" + str(data))
47. if __name__ == '__main__':
48.     generator = VulReportGenerator()
49.     generator.search_vuls()

复制代码

第四步


这时间我们可以添加其他模块功能，这里我们生成一个list模块，展示漏洞范例，并生成，-p 和-s 模块展示出分页和数量展示，便于展示,当模块多了后，我们需要设置一个-h资助模块，将内置所以功能的模块展示出来，这里每次生成的代码需要多次调试，代码有一些不足问题，需要人工辨认并指出使用ChatGpt举行辅助修改。

第五步


最后就是简报生成模块，-t，这里使用-t模块指定生成，如何生成指定到简报中？可以参考下面替换法

这个是简报模板，可以看到%S1,%S2,....,这种，分别对应着漏洞xlsx表格中的漏洞名称，风险等级，漏洞描述，加固发起，采用替换的方式，将辨认到的漏洞信息，采用替换对应的%值举行输出，另存到output文件中，为了方便也是怕导出的名称重复出现覆盖报错问题，采用递归方式001.docx，002.docx，....等方式举行下令，这里生成-t模块，将辨认到的漏洞信息指定从web.docx简报文档中举行替换%值，另存到output文件夹中的docx文件中，如许即可完成该模块需求。

到此，我们的工具思绪大致如上所述。

完整代码如下


当前版本提供两个功能1.调用本地漏洞库查询漏洞、2.批量输出漏洞简报模版

1. import argparse
2. import os
3. import pandas as pd
4. from docx import Document
6. # 创建命令行参数解析器
7. print('''
8. \033[31m  ____                              _ _            \033[0m
9. \033[31m | __ ) _   _  __ _  __      ___ __(_| |_ ___ _ __ \033[0m
10. \033[31m |  _ \| | | |/ _` | \ \ /\ / | '__| | __/ _ | '__|\033[0m     
11. \033[31m | |_) | |_| | (_| |  \ V  V /| |  | | ||  __| |   \033[0m     
12. \033[31m |____/ \__,_|\__, |   \_/\_/ |_|  |_|\__\___|_|   \033[0m     Version: \033[33m公测版本 \033[0m
13. \033[31m              |___/\033[0m                                     Tools: \033[33m漏洞报告辅助工具\033[0m      
14.             \033[34mgithub.com/yichensec/Bug_writer\033[0m     
16. 一款面向用于安服，渗透测试人员，网络安全从业者等人群的漏洞报告辅助工具可以生成漏洞测试简报，可自己私人定制。''')
17. print('''\033[33m================================== 渗透漏洞报告书写员的开始 =========================================\033[0m''')
18. class VulReportGenerator:
19.     def __init__(self):
20.         self.parser = argparse.ArgumentParser(description='\033[33m漏洞报告生成器\033[0m')
21.         self.df = pd.read_excel('./config/WEB.xlsx')
23.     def parse_args(self):
24.         self.parser.add_argument('-c', type=str, help='指定模糊查询漏洞名称.')
25.         self.parser.add_argument('-t', action="store_true", help='指定生成漏洞报告')
26.         self.parser.add_argument('-list', action='store_true', help='列出漏洞名称')
27.         self.parser.add_argument('-p', type=int, default=5, help='指定页面数量')
28.         self.parser.add_argument('-s', type=int, default=10, help='指定每页显示的项目数')
29.         self.args = self.parser.parse_args()
31.     def run(self):
32.         self.parse_args()
34.         if self.args.list:
35.             self.list_vul_names()
36.         elif self.args.c:
37.             self.search_vuls()
38.         else:
39.             print('\033[31m<<无效的参数--您可以输入"-h"查看使用说明>>\033[0m')
41.     def search_vuls(self):
42.         search_name = self.args.c.lower()
43.         result = self.df[self.df['漏洞名称'].str.lower().str.contains(search_name)]
44.         result_count = result.count()[0]
46.         if result.empty:
47.             print('\033[31m<<这里并没有发现结果--您可以去查看是否输入错误>>\033[0m')
48.         else:
49.             print(f'\033[31m查询到 {result_count} 条记录：\033[0m')
51.             total_count = result_count
52.             page_count = self.args.p
53.             rows_per_page = self.args.s
54.             current_page = 1
55.             while True:
56.                 page_result = result[(current_page-1) * rows_per_page : current_page * rows_per_page]
57.                 if page_result.empty:
58.                     break
60.                 print(f'\033[0m\033[33m第{current_page}页:\033[0m')
61.                 for row in page_result.itertuples(index=False):
62.                     print('\033[0m\033[33m漏洞名称:\033[0m\033[32m  {}\n\033[0m\033[33m风险级别:\033[0m\033[31m  {}\033[0m\n\033[33m漏洞描述:\033[0m  \n\t\033[34m  {}\033[0m\n\033[33m加固建议:\033[0m\033[35m \n\t {}\033[0m\n\n\033[33m=====================================================================================\033[0m\n'.format(row.漏洞名称, row.风险级别, row.漏洞描述, row.加固建议))
64.                 if self.args.t:
65.                     self.generate_report(page_result)
67.                 current_page += 1
68.                 if current_page > page_count:
69.                     break
71.     def list_vul_names(self):
72.         ld_names = self.df['漏洞名称']
73.         if not ld_names.empty:
74.             for index, ld_name in enumerate(ld_names):
75.                 page_num = (index + 1) // self.args.s + 1
76.                 if index % self.args.s == 0:
77.                     print(f'\033[0m\033[33m第{page_num}页:\033[0m')
78.                 print(f'{ld_name}\t', end='')
79.                 if (index + 1) % self.args.s == 0:
80.                     print()
82.             print()
83.         else:
84.             print('\033[31m<<这里并没有发现结果--您可以去查看是否输入错误>>\033[0m')
86.     def generate_report(self, page_result):
87.         doc_path = './config/web.docx'
88.         if os.path.exists(doc_path):
89.             doc_output_path = './output/001.docx'
90.             counter = 1
91.             while os.path.exists(doc_output_path):
92.                 counter += 1
93.                 doc_output_path = f'./output/{counter:03}.docx'
95.             doc = Document(doc_path)
96.             for index, row in page_result.iterrows():
97.                 for para in doc.paragraphs:
98.                     if '%S1' in para.text:
99.                         para.text = para.text.replace('%S1', row['漏洞名称'])
100.                     if '%S2' in para.text:
101.                         para.text = para.text.replace('%S2', row['风险级别'])
102.                     if '%S3' in para.text:
103.                         para.text = para.text.replace('%S3', row['漏洞描述'])
104.                     if '%S4' in para.text:
105.                         para.text = para.text.replace('%S4', row['加固建议'])
106.             try:
107.                 doc.save(doc_output_path)
108.                 print(f'\033[33m已将查询结果写入到" {doc_output_path} "文件中\033[0m')
109.             except Exception as e:
110.                 print(f'\033[31m写入文件"{doc_output_path}"时出错：{e}\033[0m')
111.         else:
112.             print(f'\033[31m指定的文件"{doc_path}"不存在\033[0m')
113. if __name__ == '__main__':
114.     generator = VulReportGenerator()
115.     generator.run()

复制代码

工具使用


1.使用下令资助

1. python main.py -h
3. 当前版本是5个参数
4.   -c C        指定模糊查询漏洞名称.
5.   -t          指定生成漏洞报告
6.   -list       列出漏洞名称
7.   -p P        指定页面数量
8.   -s S        指定每页显示的项目数

复制代码

2.漏洞查询

1. python main.py -c 查询漏洞
3. 这里-c参数是指定查询漏洞名称(可模糊查询)

复制代码

3.漏洞分页展示

1. python main.py -c 注入 -p 展示页数 -s 展示数量

复制代码

4.漏洞简报模版生成

1. python main.py -c 漏洞名 -p 展示页数 -s 展示数量 -t
3. 执行-t保存在output文件夹下的docx文件中，这里001.docx是默认名依次往后递归文件名002.docx....

复制代码

5.模版样式


打开文件夹会看到漏洞模版生成后的样式，如图下所示，（注模板可以更换自己公司的模板）

6.漏洞数量

1. python main.py -list -s 1

复制代码

   原文链接：https://www.cnblogs.com/Yichensec/p/18141979/GJ_002
  体验地址：引迈 - JNPF快速开发平台_低代码开发平台_零代码开发平台_流程设计器_表单引擎_工作流引擎_软件架构

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。