mirror of
https://github.com/Ed1s0nZ/PrivHunterAI.git
synced 2025-06-21 10:20:40 +00:00
359 lines
11 KiB
Go
359 lines
11 KiB
Go
package main
|
|
|
|
import (
|
|
"net/http"
|
|
"strconv"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/xuri/excelize/v2"
|
|
)
|
|
|
|
var Resp []Result // 数据存储在全局切片中
|
|
|
|
func Index() {
|
|
r := gin.Default()
|
|
|
|
// 提供前端静态文件服务
|
|
r.LoadHTMLFiles("index.html") // 加载前端页面
|
|
r.Static("/static", "./static") // 为前端静态资源提供服务
|
|
|
|
r.GET("/", func(c *gin.Context) {
|
|
c.HTML(http.StatusOK, "index.html", nil)
|
|
})
|
|
|
|
// 分页数据接口
|
|
r.GET("/data", func(c *gin.Context) {
|
|
// 获取分页参数
|
|
page, err1 := strconv.Atoi(c.Query("page"))
|
|
pageSize, err2 := strconv.Atoi(c.Query("pageSize"))
|
|
resultFilter := c.Query("result")
|
|
|
|
if err1 != nil || page < 1 {
|
|
page = 1
|
|
}
|
|
if err2 != nil || pageSize < 1 {
|
|
pageSize = 10
|
|
}
|
|
|
|
// 应用筛选条件
|
|
var filteredData []Result
|
|
for _, item := range Resp {
|
|
if resultFilter == "" || item.Result == resultFilter {
|
|
filteredData = append(filteredData, item)
|
|
}
|
|
}
|
|
|
|
// 计算分页数据
|
|
total := len(filteredData)
|
|
offset := (page - 1) * pageSize
|
|
var data []Result
|
|
if offset < total {
|
|
if offset+pageSize > total {
|
|
data = filteredData[offset:]
|
|
} else {
|
|
data = filteredData[offset : offset+pageSize]
|
|
}
|
|
}
|
|
|
|
// 返回响应
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"data": data,
|
|
"total": total,
|
|
"currentPage": page,
|
|
"pageSize": pageSize,
|
|
"totalPages": (total + pageSize - 1) / pageSize,
|
|
})
|
|
})
|
|
|
|
// 统计数据接口
|
|
r.GET("/stats", func(c *gin.Context) {
|
|
total := len(Resp)
|
|
vulnerable := 0
|
|
unknown := 0
|
|
safe := 0
|
|
|
|
for _, item := range Resp {
|
|
switch item.Result {
|
|
case "true":
|
|
vulnerable++
|
|
case "unknown":
|
|
unknown++
|
|
case "false":
|
|
safe++
|
|
}
|
|
}
|
|
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"total": total,
|
|
"vulnerable": vulnerable,
|
|
"unknown": unknown,
|
|
"safe": safe,
|
|
})
|
|
})
|
|
|
|
// 添加数据接口
|
|
r.POST("/update", func(c *gin.Context) {
|
|
var newData Result
|
|
if err := c.ShouldBindJSON(&newData); err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
Resp = append(Resp, newData)
|
|
c.JSON(http.StatusOK, gin.H{"message": "Data updated successfully"})
|
|
})
|
|
|
|
// 导出Excel接口
|
|
r.GET("/export", func(c *gin.Context) {
|
|
resultFilter := c.Query("result")
|
|
|
|
// 应用筛选条件
|
|
var filteredData []Result
|
|
for _, item := range Resp {
|
|
if resultFilter == "" || item.Result == resultFilter {
|
|
filteredData = append(filteredData, item)
|
|
}
|
|
}
|
|
|
|
// 创建Excel文件
|
|
f := excelize.NewFile()
|
|
defer func() {
|
|
if err := f.Close(); err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
|
}
|
|
}()
|
|
|
|
// 创建基础单元格样式（仅包含边框）
|
|
baseCellStyle, _ := f.NewStyle(&excelize.Style{
|
|
Border: []excelize.Border{
|
|
{Type: "left", Color: "#000000", Style: 1},
|
|
{Type: "top", Color: "#000000", Style: 1},
|
|
{Type: "bottom", Color: "#000000", Style: 1},
|
|
{Type: "right", Color: "#000000", Style: 1},
|
|
},
|
|
})
|
|
|
|
// 创建漏洞状态的单元格样式
|
|
vulnerableStyle, _ := f.NewStyle(&excelize.Style{
|
|
Border: []excelize.Border{
|
|
{Type: "left", Color: "#000000", Style: 1},
|
|
{Type: "top", Color: "#000000", Style: 1},
|
|
{Type: "bottom", Color: "#000000", Style: 1},
|
|
{Type: "right", Color: "#000000", Style: 1},
|
|
},
|
|
Fill: excelize.Fill{
|
|
Type: "pattern",
|
|
Color: []string{"#ffebee"},
|
|
Pattern: 1,
|
|
},
|
|
})
|
|
|
|
// 创建未知状态的单元格样式
|
|
unknownStyle, _ := f.NewStyle(&excelize.Style{
|
|
Border: []excelize.Border{
|
|
{Type: "left", Color: "#000000", Style: 1},
|
|
{Type: "top", Color: "#000000", Style: 1},
|
|
{Type: "bottom", Color: "#000000", Style: 1},
|
|
{Type: "right", Color: "#000000", Style: 1},
|
|
},
|
|
Fill: excelize.Fill{
|
|
Type: "pattern",
|
|
Color: []string{"#fff8e1"},
|
|
Pattern: 1,
|
|
},
|
|
})
|
|
|
|
// 创建安全状态的单元格样式
|
|
safeStyle, _ := f.NewStyle(&excelize.Style{
|
|
Border: []excelize.Border{
|
|
{Type: "left", Color: "#000000", Style: 1},
|
|
{Type: "top", Color: "#000000", Style: 1},
|
|
{Type: "bottom", Color: "#000000", Style: 1},
|
|
{Type: "right", Color: "#000000", Style: 1},
|
|
},
|
|
Fill: excelize.Fill{
|
|
Type: "pattern",
|
|
Color: []string{"#e8f5e9"},
|
|
Pattern: 1,
|
|
},
|
|
})
|
|
|
|
// 设置单元格样式
|
|
headerStyle, _ := f.NewStyle(&excelize.Style{
|
|
Font: &excelize.Font{
|
|
Bold: true,
|
|
Color: "#FFFFFF",
|
|
},
|
|
Fill: excelize.Fill{
|
|
Type: "pattern",
|
|
Color: []string{"#4a90e2"},
|
|
Pattern: 1,
|
|
},
|
|
Border: []excelize.Border{
|
|
{Type: "left", Color: "#000000", Style: 1},
|
|
{Type: "top", Color: "#000000", Style: 1},
|
|
{Type: "bottom", Color: "#000000", Style: 1},
|
|
{Type: "right", Color: "#000000", Style: 1},
|
|
},
|
|
Alignment: &excelize.Alignment{
|
|
Horizontal: "center",
|
|
Vertical: "center",
|
|
},
|
|
})
|
|
|
|
// 创建漏洞状态的文本换行样式
|
|
vulnerableWrapStyle, _ := f.NewStyle(&excelize.Style{
|
|
Alignment: &excelize.Alignment{
|
|
WrapText: true,
|
|
Vertical: "top",
|
|
},
|
|
Border: []excelize.Border{
|
|
{Type: "left", Color: "#000000", Style: 1},
|
|
{Type: "top", Color: "#000000", Style: 1},
|
|
{Type: "bottom", Color: "#000000", Style: 1},
|
|
{Type: "right", Color: "#000000", Style: 1},
|
|
},
|
|
Fill: excelize.Fill{
|
|
Type: "pattern",
|
|
Color: []string{"#fff0f0"},
|
|
Pattern: 1,
|
|
},
|
|
})
|
|
|
|
// 创建未知状态的文本换行样式
|
|
unknownWrapStyle, _ := f.NewStyle(&excelize.Style{
|
|
Alignment: &excelize.Alignment{
|
|
WrapText: true,
|
|
Vertical: "top",
|
|
},
|
|
Border: []excelize.Border{
|
|
{Type: "left", Color: "#000000", Style: 1},
|
|
{Type: "top", Color: "#000000", Style: 1},
|
|
{Type: "bottom", Color: "#000000", Style: 1},
|
|
{Type: "right", Color: "#000000", Style: 1},
|
|
},
|
|
Fill: excelize.Fill{
|
|
Type: "pattern",
|
|
Color: []string{"#fffaec"},
|
|
Pattern: 1,
|
|
},
|
|
})
|
|
|
|
// 创建安全状态的文本换行样式
|
|
safeWrapStyle, _ := f.NewStyle(&excelize.Style{
|
|
Alignment: &excelize.Alignment{
|
|
WrapText: true,
|
|
Vertical: "top",
|
|
},
|
|
Border: []excelize.Border{
|
|
{Type: "left", Color: "#000000", Style: 1},
|
|
{Type: "top", Color: "#000000", Style: 1},
|
|
{Type: "bottom", Color: "#000000", Style: 1},
|
|
{Type: "right", Color: "#000000", Style: 1},
|
|
},
|
|
Fill: excelize.Fill{
|
|
Type: "pattern",
|
|
Color: []string{"#f0fff0"},
|
|
Pattern: 1,
|
|
},
|
|
})
|
|
|
|
// 设置表头
|
|
sheetName := "扫描结果"
|
|
f.SetSheetName("Sheet1", sheetName)
|
|
headers := []string{"API地址", "请求方式", "状态", "置信度", "原因", "原始请求", "原始响应", "重放请求", "重放响应", "时间戳"}
|
|
for i, header := range headers {
|
|
col := string(rune('A' + i))
|
|
f.SetCellValue(sheetName, col+"1", header)
|
|
f.SetCellStyle(sheetName, col+"1", col+"1", headerStyle)
|
|
}
|
|
|
|
// 设置列宽
|
|
f.SetColWidth(sheetName, "A", "A", 50) // API地址
|
|
f.SetColWidth(sheetName, "B", "B", 10) // 请求方式
|
|
f.SetColWidth(sheetName, "C", "C", 15) // 状态
|
|
f.SetColWidth(sheetName, "D", "D", 15) // 置信度
|
|
f.SetColWidth(sheetName, "E", "E", 40) // 原因
|
|
f.SetColWidth(sheetName, "F", "F", 50) // 原始请求
|
|
f.SetColWidth(sheetName, "G", "G", 50) // 原始响应
|
|
f.SetColWidth(sheetName, "H", "H", 50) // 重放请求
|
|
f.SetColWidth(sheetName, "I", "I", 50) // 重放响应
|
|
f.SetColWidth(sheetName, "J", "J", 20) // 时间戳
|
|
|
|
// 填充数据
|
|
for i, item := range filteredData {
|
|
row := i + 2 // 从第2行开始(表头是第1行)
|
|
|
|
// 状态显示文本与样式
|
|
var statusText string
|
|
var styleID int
|
|
if item.Result == "true" {
|
|
statusText = "漏洞"
|
|
styleID = vulnerableStyle
|
|
} else if item.Result == "unknown" {
|
|
statusText = "未知"
|
|
styleID = unknownStyle
|
|
} else {
|
|
statusText = "安全"
|
|
styleID = safeStyle
|
|
}
|
|
|
|
f.SetCellValue(sheetName, "A"+strconv.Itoa(row), item.Url)
|
|
f.SetCellValue(sheetName, "B"+strconv.Itoa(row), item.Method)
|
|
f.SetCellValue(sheetName, "C"+strconv.Itoa(row), statusText)
|
|
f.SetCellValue(sheetName, "D"+strconv.Itoa(row), item.Confidence)
|
|
f.SetCellValue(sheetName, "E"+strconv.Itoa(row), item.Reason)
|
|
f.SetCellValue(sheetName, "F"+strconv.Itoa(row), item.RequestA)
|
|
f.SetCellValue(sheetName, "G"+strconv.Itoa(row), item.RespBodyA)
|
|
f.SetCellValue(sheetName, "H"+strconv.Itoa(row), item.RequestB)
|
|
f.SetCellValue(sheetName, "I"+strconv.Itoa(row), item.RespBodyB)
|
|
f.SetCellValue(sheetName, "J"+strconv.Itoa(row), item.Timestamp)
|
|
|
|
// 设置普通单元格的基础样式
|
|
f.SetCellStyle(sheetName, "A"+strconv.Itoa(row), "A"+strconv.Itoa(row), baseCellStyle)
|
|
f.SetCellStyle(sheetName, "B"+strconv.Itoa(row), "B"+strconv.Itoa(row), baseCellStyle)
|
|
f.SetCellStyle(sheetName, "C"+strconv.Itoa(row), "C"+strconv.Itoa(row), styleID)
|
|
f.SetCellStyle(sheetName, "D"+strconv.Itoa(row), "D"+strconv.Itoa(row), baseCellStyle)
|
|
f.SetCellStyle(sheetName, "J"+strconv.Itoa(row), "J"+strconv.Itoa(row), baseCellStyle)
|
|
|
|
// 根据状态，为文本换行的单元格设置对应的背景色
|
|
if item.Result == "true" {
|
|
f.SetCellStyle(sheetName, "E"+strconv.Itoa(row), "E"+strconv.Itoa(row), vulnerableWrapStyle)
|
|
f.SetCellStyle(sheetName, "F"+strconv.Itoa(row), "F"+strconv.Itoa(row), vulnerableWrapStyle)
|
|
f.SetCellStyle(sheetName, "G"+strconv.Itoa(row), "G"+strconv.Itoa(row), vulnerableWrapStyle)
|
|
f.SetCellStyle(sheetName, "H"+strconv.Itoa(row), "H"+strconv.Itoa(row), vulnerableWrapStyle)
|
|
f.SetCellStyle(sheetName, "I"+strconv.Itoa(row), "I"+strconv.Itoa(row), vulnerableWrapStyle)
|
|
} else if item.Result == "unknown" {
|
|
f.SetCellStyle(sheetName, "E"+strconv.Itoa(row), "E"+strconv.Itoa(row), unknownWrapStyle)
|
|
f.SetCellStyle(sheetName, "F"+strconv.Itoa(row), "F"+strconv.Itoa(row), unknownWrapStyle)
|
|
f.SetCellStyle(sheetName, "G"+strconv.Itoa(row), "G"+strconv.Itoa(row), unknownWrapStyle)
|
|
f.SetCellStyle(sheetName, "H"+strconv.Itoa(row), "H"+strconv.Itoa(row), unknownWrapStyle)
|
|
f.SetCellStyle(sheetName, "I"+strconv.Itoa(row), "I"+strconv.Itoa(row), unknownWrapStyle)
|
|
} else {
|
|
f.SetCellStyle(sheetName, "E"+strconv.Itoa(row), "E"+strconv.Itoa(row), safeWrapStyle)
|
|
f.SetCellStyle(sheetName, "F"+strconv.Itoa(row), "F"+strconv.Itoa(row), safeWrapStyle)
|
|
f.SetCellStyle(sheetName, "G"+strconv.Itoa(row), "G"+strconv.Itoa(row), safeWrapStyle)
|
|
f.SetCellStyle(sheetName, "H"+strconv.Itoa(row), "H"+strconv.Itoa(row), safeWrapStyle)
|
|
f.SetCellStyle(sheetName, "I"+strconv.Itoa(row), "I"+strconv.Itoa(row), safeWrapStyle)
|
|
}
|
|
|
|
// 设置行高，确保内容能够正常显示
|
|
f.SetRowHeight(sheetName, row, 100)
|
|
}
|
|
|
|
// 生成Excel文件
|
|
c.Header("Content-Type", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet")
|
|
c.Header("Content-Disposition", "attachment; filename=PrivHunterAI-扫描结果-"+time.Now().Format("2006-01-02-15-04-05")+".xlsx")
|
|
c.Header("Content-Transfer-Encoding", "binary")
|
|
|
|
// 将Excel文件写入响应
|
|
if err := f.Write(c.Writer); err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
})
|
|
|
|
// 启动服务
|
|
r.Run(":8222")
|
|
}
|