最近也開通了 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 是不可能對整個源碼進行全文解析的，只能進行小段解析。
所以我想用手段先進行預處理。用一些軟件的漏洞匹配規則進行預掃描可以確定出某一行代碼存在風險隱患。然後針對該行代碼有兩種操作，判斷其是否在函數內。如果是函數則函數整體分析，如果不是則在一個閾值內進行上下文聯動分析。
高級點的玩法還可以跟蹤變量賦值操作。這樣基本已經完成整個程序審計邏輯了。
其實可以玩的方向還有很多，這裡就不再多說。
代碼其實已經寫完現在捂在手裡，如果四月過完公司沒有催上交我就把它開源出來。
小小的露一個圖
最後
我不經常寫文章所以文筆不好語言混亂，望大家見諒。本文提及的一些東西算是比較熱門的話題，請大家以純技術的方向看待，擁抱開源但不要做白嫖賊