基于GPT的网络安全工具设计

最近也开通了 GPT puls，感觉 GPT4 的效果确实比 3.5 强很多。不仅仅是上下文能力和字数限制，逻辑能力也提高不不少。终于可以将其用在生产力解放上面，于是便有了这么一篇文章。不是专业开发，代码稀烂，大佬轻喷

01#

众所周知安全岗位工作的很大一部分工作就是撰写报告。往往一个漏洞用五分钟发现，书写报告却用了十五分钟。一篇报告要按照特定的格式特定的语法描述进行书写，实在是一件让人煎熬的事，那么，我们就需要让 AI 把我们从痛苦中解放出来。

设计思路#

拥有一个 UI 界面
便于文件处理
基于 AI 处理语言描述问题
可以人工干预结果

首先写一个 ui 界面。为了快捷开发我选择了 python，所以直接选择了 web 进行轻便开发。

flask 起手

app = Flask(__name__, template_folder='templates', static_folder='static')

@app.route('/', methods=['GET', 'POST'])
def index():
    if request.method == 'POST':
        vuln_name = request.form.get('vuln_name', default="SQL注入")
        vuln_point = request.form.get('vuln_point', default="www.google.com")
        beizhu = request.form.get('beizhu', default="")
        language = request.form.get('language')
        start = time.time()
        report = generate_report(vuln_name, vuln_point, beizhu, language)
        end = time.time()
        times = end - start

        return render_template('index.html', report=report, times=times)
    else:
        return render_template('index.html')

创建了一个默认路由，并且加载模板 index.html
index.html 主体是一个表单。

表单：

漏洞名称（eg. SQL 注入）
漏洞点（eg. www.test.com/id=111 ）
备注（比如报告详细些、给出 xx 个 POC 之类的额外要求）

那么再写一个 html 来简单的实现这个页面

<!DOCTYPE html>
<html>

<head>
  <meta charset="UTF-8">
  <title>漏洞报告生成器</title>
  <style>
    body {
      font-family: Arial, sans-serif;
      margin: 0;
      padding: 0;
    }

    header {
      background-color: #005293;
      color: #fff;
      padding: 10px;
    }

    h1 {
      margin: 0;
    }

    form {
      margin: 20px;
    }

    label {
      display: block;
      margin-bottom: 10px;
    }

    input[type="text"] {
      width: 100%;
      padding: 10px;
      border: 1px solid #ccc;
      border-radius: 5px;
      margin-bottom: 20px;
      font-size: 16px;
    }

    input[type="submit"] {
      background-color: #005293;
      color: #fff;
      border: none;
      padding: 10px;
      border-radius: 5px;
      cursor: pointer;
      font-size: 16px;
    }

    #report {
      margin: 20px;
      border: 1px solid #ccc;
      padding: 10px;
      border-radius: 5px;
      font-size: 16px;
    }
  </style>

  <link rel="stylesheet"
    href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.7.0/styles/panda-syntax-dark.min.css">

  <script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.7.0/highlight.min.js"
    integrity="sha512-bgHRAiTjGrzHzLyKOnpFvaEpGzJet3z4tZnXGjpsCcqOnAH6VGUx9frc5bcIhKTVLEiCO6vEhNAgx5jtLUYrfA=="
    crossorigin="anonymous" referrerpolicy="no-referrer"></script>

  <script>hljs.initHighlightingOnLoad();</script>
  <script src="https://cdn.jsdelivr.net/npm/[email protected]/base64.min.js"></script>
</head>

<body>
  <header>
    <h1>漏洞报告生成器</h1>

  </header>

  <form method="POST" action="/">

    <select name="language">
      <option value="中文">中文</option>
      <option value="英文">English</option>
    </select>

    <label>漏洞名称：</label>
    <input type="text" name="vuln_name" placeholder="请输入漏洞名称">
    <label>漏洞点：</label>
    <input type="text" name="vuln_point" placeholder="请输入漏洞点">
    <label>对报告的额外要求：</label>
    <input type="text" name="beizhu" placeholder="报告额外要求">

    <input type="submit" value="生成漏洞报告"> <button id="copy-btn" onclick="copyReport()" type="button">复制</button><!--<button id="copy-btn" onclick="saveReport()" type="button">生成文件</button>-->
    本次执行耗时：{{ times }} 秒
  </form>


  {% if report %}

  <pre id="report" style="white-space: pre-wrap;overflow-wrap: break-word;">
        <code>
      {{ report }}
      </code>
    </pre>


  <br>

  {% if report %}

  {% endif %}

  <script>
    function copyReport() {
      var reportText = document.querySelector("#report code").innerText;
      navigator.clipboard.writeText(reportText)
        .then(function () {
          alert("漏洞报告已复制到剪贴板！");
        })
        .catch(function (error) {
          alert("复制失败，请手动复制。");
        });
    }
  </script>

  {% endif %}
</body>

</html>

页面大概就这个样，顺便完善一些小细节。

然后是文件处理。有没有一种文件或者编码形式可以让人愉快的阅读而且可以被浏览器文本编辑器解析呢？没错，makedown 首当其冲。于是在 index.html 里加上了代码高亮，生成的值直接填充到其中。

最终实现效果：

然后是人工调控，但其实已经由备注选项完成，不再论述。

最后是最关键的

对接 chatGPT api 并且获取信息
这里推荐一个 github 上的项目，关于 chatGPT 的开发资料可以在这里查到：https://github.com/easychen/openai-gpt-dev-notes-for-cn-developer

接下来写逻辑代码

def generate_report(vuln_name, vuln_point, beihzu="", language=""):
    # 构造 GPT-3 输入
    prompt = f"" #这里填写调整好的prompt

    #api_base = {"socks5":proxy}
    api_base = {"http": proxy, "https": proxy}

    # 调用 OpenAI API 生成漏洞报告
    headers = {
        # Already added when you pass json= but not when you pass data=
        'Content-Type': 'application/json',
        'Authorization': "Bearer "+api_key,
    }

    json_data = {
        'model': 'gpt-3.5-turbo',
        'messages': [
            {
                'role': 'user',
                'content': prompt,
            },
        ],
    }
    # 设置代理
    # 计算时间

    response = requests.post('https://api.openai.com/v1/chat/completions',
                             headers=headers, json=json_data, proxies=api_base)

    # 处理 OpenAI API 的响应
    if response.status_code == 200:
        text = response.text
        text = json.loads(text)
        text = text['choices'][0]['message']['content']
        report = f"{text}"

    else:
        report = "生成漏洞报告失败，请检查输入并重试。状态码："+str(response.status_code)

    return report

这里没有用官方的方法进行调用，而是为了方便代理用了 http 请求接口。

大致逻辑就是这样，顺便再写一个文件保存功能

def outfile():
    #获取post过来的base64编码的date
    date=request.form.get('reportText')
    
    #将date进行base64解码
    date=base64.b64decode(date,).decode('utf-8')
    #将date写入md文件，文件名为当前时间
    fname=time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
    with open(fname+'.md','w') as f:
        f.write(date)
    return "ok"

总体就是这样，最终的效果图

最后缝合一下，项目放在了
https://github.com/shiyeshu/GPTreport

大佬们喜欢的可以给个 star

02#

本来其实这个才是重头戏，但是公司有一个莫名其妙的要求。要让四月底注册一个专利献祭给公司。搜肠刮肚没办法准备把这个献祭掉。

大致说说想法

做一个 AI 的代码审计工具，肯定是比前面的复杂且更实用。

首先 ai 是不可能对整个源码进行全文解析的，只能进行小段解析。
所以我想用手段先进行预处理。用一些软件的漏洞匹配规则进行预扫描可以确定出某一行代码存在风险隐患。然后针对该行代码有两种操作，判断其是否在函数内。如果是函数则函数整体分析，如果不是则在一个阈值内进行上下文联动分析。
高级点的玩法还可以跟踪变量赋值操作。这样基本已经完成整个程序审计逻辑了。
其实可以玩的方向还有很多，这里就不再多说。
代码其实已经写完现在捂在手里，如果四月过完公司没有催上交我就把它开源出来。
小小的露一个图

最后
我不经常写文章所以文笔不好语言混乱，望大家见谅。本文提及的一些东西算是比较热门的话题，请大家以纯技术的方向看待，拥抱开源但不要做白嫖贼