diff --git a/index.go b/index.go index cff39a0..32ef1a2 100644 --- a/index.go +++ b/index.go @@ -3,8 +3,10 @@ package main import ( "net/http" "strconv" + "time" "github.com/gin-gonic/gin" + "github.com/xuri/excelize/v2" ) var Resp []Result // 数据存储在全局切片中 @@ -101,6 +103,256 @@ func Index() { 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") }