写点什么

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

  • 2024-04-18
    福建
  • 本文字数:6656 字

    阅读完需:约 22 分钟

前言


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


案例一 漏洞简报生成


工具前景


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


工具名称


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


工具结构


工具结构Bug_writer:│  main.py│  README.md│  requirements.txt├─config│  │  WEB.docx│  │  WEB.xlsx│  ││  └─漏洞模版│          简报模版一.docx│          默认模板.docx└─output 说明: 主目录 mian.py为脚本文件│  README.md 为使用说明│  requirements.txt 为需要安装的第三方库│  config文件夹下│    WEB.docx 简报模版│    WEB.xlsx 漏洞库│    漏洞模版文件夹下,为模版样式│  output文件夹为输出的简报文件
复制代码


工具思路


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



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



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


第一步


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




第二步


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



代码如下

import argparseimport osimport pandas as pd  class VulReportGenerator:    def __init__(self):        self.parser = argparse.ArgumentParser(description='漏洞报告生成器')        self.parser.add_argument('-c', type=str, help='指定模糊查询漏洞名称.')        self.parser.add_argument('-t', action='store_true', help='生成报告')        self.args = self.parser.parse_args()        self.df = pd.read_excel('./config/WEB.xlsx')     def search_vuls(self):        search_name = self.args.c.lower()        result = self.df[self.df['漏洞名称'].str.lower().str.contains(search_name)]        result_count = result.shape[0]         if result.empty:            print('<<这里并没有发现结果--您可以去查看是否输入错误>>')        else:            print(f'查询到 {result_count} 条记录:')             for index, row in result.iterrows():                print('漏洞名称:  {}\n风险级别: {}\n漏洞描述:\n\t {} \n加固建议: \n\t {}\n\n\n'.format(                    row['漏洞名称'], row['风险级别'], row['漏洞描述'], row['加固建议']))             if self.args.t:                self.generate_report(result)     def generate_report(self, data):        # 在这里编写生成报告的逻辑,例如将结果输出到文件或者其他操作        print("生成报告:" + str(data))  if __name__ == '__main__':    generator = VulReportGenerator()    generator.search_vuls()
复制代码


第三步


这一步的时候,我们可以看到,样式有些难看,我们可以使用图形化字符工具生成一个字符图形,然后再完善增加颜色样式。




代码如下

import argparseimport osimport pandas as pd # 创建命令行参数解析器print(''' \033[31m  ____                              _ _            \033[0m \033[31m | __ ) _   _  __ _  __      ___ __(_| |_ ___ _ __ \033[0m \033[31m |  _ \| | | |/ _` | \ \ /\ / | '__| | __/ _ | '__|\033[0m      \033[31m | |_) | |_| | (_| |  \ V  V /| |  | | ||  __| |   \033[0m      \033[31m |____/ \__,_|\__, |   \_/\_/ |_|  |_|\__\___|_|   \033[0m     Version: \033[33m公测版本 \033[0m \033[31m              |___/\033[0m                                     Tools: \033[33m漏洞报告辅助工具\033[0m                  \033[34mgithub.com/yichensec/Bug_writer\033[0m      一款面向用于安服,渗透测试人员,网络安全从业者等人群的漏洞报告辅助工具可以生成漏洞测试简报,可自己私人定制。''')print('''\033[33m================================== 渗透漏洞报告书写员的开始 =========================================\033[0m''') class VulReportGenerator:    def __init__(self):        self.parser = argparse.ArgumentParser(description='漏洞报告生成器')        self.parser.add_argument('-c', type=str, help='指定模糊查询漏洞名称.')        self.parser.add_argument('-t', action='store_true', help='生成报告')        self.args = self.parser.parse_args()        self.df = pd.read_excel('./config/WEB.xlsx')     def search_vuls(self):        search_name = self.args.c.lower()        result = self.df[self.df['漏洞名称'].str.lower().str.contains(search_name)]        result_count = result.shape[0]         if result.empty:            print('<<这里并没有发现结果--您可以去查看是否输入错误>>')        else:            print(f'查询到 {result_count} 条记录:')             for index, row in result.iterrows():                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.加固建议))             if self.args.t:                self.generate_report(result)     def generate_report(self, data):        # 在这里编写生成报告的逻辑,例如将结果输出到文件或者其他操作        print("生成报告:" + str(data))  if __name__ == '__main__':    generator = VulReportGenerator()    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.批量输出漏洞简报模版

import argparseimport osimport pandas as pdfrom docx import Document # 创建命令行参数解析器print(''' \033[31m  ____                              _ _            \033[0m \033[31m | __ ) _   _  __ _  __      ___ __(_| |_ ___ _ __ \033[0m \033[31m |  _ \| | | |/ _` | \ \ /\ / | '__| | __/ _ | '__|\033[0m      \033[31m | |_) | |_| | (_| |  \ V  V /| |  | | ||  __| |   \033[0m      \033[31m |____/ \__,_|\__, |   \_/\_/ |_|  |_|\__\___|_|   \033[0m     Version: \033[33m公测版本 \033[0m \033[31m              |___/\033[0m                                     Tools: \033[33m漏洞报告辅助工具\033[0m                  \033[34mgithub.com/yichensec/Bug_writer\033[0m      一款面向用于安服,渗透测试人员,网络安全从业者等人群的漏洞报告辅助工具可以生成漏洞测试简报,可自己私人定制。''')print('''\033[33m================================== 渗透漏洞报告书写员的开始 =========================================\033[0m''')class VulReportGenerator:    def __init__(self):        self.parser = argparse.ArgumentParser(description='\033[33m漏洞报告生成器\033[0m')        self.df = pd.read_excel('./config/WEB.xlsx')     def parse_args(self):        self.parser.add_argument('-c', type=str, help='指定模糊查询漏洞名称.')        self.parser.add_argument('-t', action="store_true", help='指定生成漏洞报告')        self.parser.add_argument('-list', action='store_true', help='列出漏洞名称')        self.parser.add_argument('-p', type=int, default=5, help='指定页面数量')        self.parser.add_argument('-s', type=int, default=10, help='指定每页显示的项目数')        self.args = self.parser.parse_args()     def run(self):        self.parse_args()         if self.args.list:            self.list_vul_names()        elif self.args.c:            self.search_vuls()        else:            print('\033[31m<<无效的参数--您可以输入"-h"查看使用说明>>\033[0m')     def search_vuls(self):        search_name = self.args.c.lower()        result = self.df[self.df['漏洞名称'].str.lower().str.contains(search_name)]        result_count = result.count()[0]         if result.empty:            print('\033[31m<<这里并没有发现结果--您可以去查看是否输入错误>>\033[0m')        else:            print(f'\033[31m查询到 {result_count} 条记录:\033[0m')             total_count = result_count            page_count = self.args.p            rows_per_page = self.args.s            current_page = 1            while True:                page_result = result[(current_page-1) * rows_per_page : current_page * rows_per_page]                if page_result.empty:                    break                 print(f'\033[0m\033[33m第{current_page}页:\033[0m')                for row in page_result.itertuples(index=False):                    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.加固建议))                 if self.args.t:                    self.generate_report(page_result)                 current_page += 1                if current_page > page_count:                    break     def list_vul_names(self):        ld_names = self.df['漏洞名称']        if not ld_names.empty:            for index, ld_name in enumerate(ld_names):                page_num = (index + 1) // self.args.s + 1                if index % self.args.s == 0:                    print(f'\033[0m\033[33m第{page_num}页:\033[0m')                print(f'{ld_name}\t', end='')                if (index + 1) % self.args.s == 0:                    print()             print()        else:            print('\033[31m<<这里并没有发现结果--您可以去查看是否输入错误>>\033[0m')     def generate_report(self, page_result):        doc_path = './config/web.docx'        if os.path.exists(doc_path):            doc_output_path = './output/001.docx'            counter = 1            while os.path.exists(doc_output_path):                counter += 1                doc_output_path = f'./output/{counter:03}.docx'             doc = Document(doc_path)            for index, row in page_result.iterrows():                for para in doc.paragraphs:                    if '%S1' in para.text:                        para.text = para.text.replace('%S1', row['漏洞名称'])                    if '%S2' in para.text:                        para.text = para.text.replace('%S2', row['风险级别'])                    if '%S3' in para.text:                        para.text = para.text.replace('%S3', row['漏洞描述'])                    if '%S4' in para.text:                        para.text = para.text.replace('%S4', row['加固建议'])            try:                doc.save(doc_output_path)                print(f'\033[33m已将查询结果写入到" {doc_output_path} "文件中\033[0m')            except Exception as e:                print(f'\033[31m写入文件"{doc_output_path}"时出错:{e}\033[0m')        else:            print(f'\033[31m指定的文件"{doc_path}"不存在\033[0m')if __name__ == '__main__':    generator = VulReportGenerator()    generator.run()
复制代码


工具使用


1.使用命令帮助

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



2.漏洞查询

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



3.漏洞分页展示


python main.py -c 注入 -p 展示页数 -s 展示数量
复制代码



4.漏洞简报模版生成


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



5.模版样式


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



6.漏洞数量


python main.py -list -s 1
复制代码




原文链接:https://www.cnblogs.com/Yichensec/p/18141979/GJ_002

体验地址:http://www.jnpfsoft.com/?from=001

用户头像

还未添加个人签名 2023-06-19 加入

还未添加个人简介

评论

发布
暂无评论
基于ChatGPT打造安全脚本工具流程_安全_不在线第一只蜗牛_InfoQ写作社区