添加deepseek支持

common/runner/runner.go

@@ -4,6 +4,7 @@ package runner
 import (
 	"bufio"
 	"fmt"
+	"github.com/Tencent/AI-Infra-Guard/pkg/openai"
 	"math"
 	"net/http"
 	"os"
@@ -18,7 +19,6 @@ import (
 	"github.com/Tencent/AI-Infra-Guard/internal/gologger"
 	"github.com/Tencent/AI-Infra-Guard/internal/options"
 	"github.com/Tencent/AI-Infra-Guard/pkg/httpx"
-	"github.com/Tencent/AI-Infra-Guard/pkg/hunyuan"
 	"github.com/Tencent/AI-Infra-Guard/pkg/vulstruct"
 
 	"github.com/liushuochen/gotable"
@@ -559,13 +559,19 @@ func (r *Runner) writeResult(f *os.File, result HttpResult) {
 		}
 		if r.Options.AIAnalysis {
 			fmt.Println("AI分析:")
-			prompt := "你是安全漏洞报告解读大师，我会给你扫描器输出的url和存在的cve详情。以编写甲方漏洞报告的形式编写完整报告，参考格式如：\n# 一、风险总览\n(描述测试的url以及基本信息，综合CVE漏洞可能造成的严重漏洞后果)\n# 二、漏洞详情\n(请你利用搜索等功能，依次分析CVE的详情，给出漏洞怎么产生，怎么利用，修复方案的详情(根据漏洞类型给出对应修复方案而不是简单升级)，然后给出可靠参考来源,相同类型漏洞合并在一起给出)\n漏洞报告如下：\n"
+			prompt := "你是安全漏洞报告解读大师，我会给你扫描器输出的url和存在的cve详情。以编写甲方漏洞报告的形式编写完整报告，参考格式如：\n# 一、风险总览\n(描述测试的url以及基本信息，综合CVE漏洞可能造成的严重漏洞后果)\n# 二、漏洞详情\n(请你利用搜索等功能，依次分析CVE的详情，给出漏洞怎么产生，怎么利用，修复方案的详情(根据漏洞类型给出对应修复方案，执行的命令,而不是简单升级)，然后给出可靠参考来源,相同类型漏洞合并在一起给出)\n漏洞报告如下：\n"
 			prompt += fmt.Sprintf("%s title:%s fingerprint:%v", result.URL, result.Title, result.Fingers) + "\n"
 			for _, item := range result.Advisories {
 				prompt += fmt.Sprintf("%s[%s]:%s\n", item.Info.CVEName, item.Info.Severity, item.Info.Details)
 				prompt += fmt.Sprintf("reference: %v\n", item.References)
 			}
-			full, err := hunyuan.HunyuanAI(prompt, r.Options.AIToken)
+			var err error
+			var full string
+			if r.Options.AIDeepSeekToken != "" {
+				full, err = openai.DeepSeekR1API(prompt, r.Options.AIDeepSeekToken)
+			} else {
+				full, err = openai.HunyuanAI(prompt, r.Options.AIHunyuanToken)
+			}
 			if err != nil {
 				gologger.WithError(err).Errorln("AI分析失败")
 			}

internal/options/banner.go

@@ -6,7 +6,7 @@ import (
 	"fmt"
 )
 
-const version = "v0.0.6-dev"
+const version = "v0.0.7-dev"
 
 // ShowBanner is used to show the banner to the user
 func ShowBanner() {

internal/options/options.go

@@ -23,7 +23,8 @@ type Options struct {
 	ListVulTemplate bool              // 是否列出漏洞模板
 	CheckVulTargets bool              // 检查漏洞模板是否正确
 	AIAnalysis      bool              // 是否启用AI分析
-	AIToken         string            // AI服务的认证令牌
+	AIHunyuanToken  string            // AI服务的认证令牌
+	AIDeepSeekToken string            // deepseek的认证令牌
 	LocalScan       bool              // 一键检测本地
 	WebServer       bool              // 是否启用WebSocket服务器
 	WebServerAddr   string            // WebSocket服务器地址
@@ -59,11 +60,12 @@ func ParseOptions() *Options {
 	flag.StringVar(&options.AdvTemplates, "vul", "data/vuln", "指定漏洞库目录")
 	flag.BoolVar(&options.ListVulTemplate, "list-vul", false, "输出漏洞库列表")
 	flag.BoolVar(&options.CheckVulTargets, "check-vul", false, "检查漏洞模板是否正确")
-	flag.BoolVar(&options.AIAnalysis, "ai", false, "AI分析")
-	flag.StringVar(&options.AIToken, "token", "", "混元token")
 	flag.BoolVar(&options.LocalScan, "localscan", false, "一键检测本地")
 	flag.BoolVar(&options.WebServer, "ws", false, "启用WebServer服务器")
 	flag.StringVar(&options.WebServerAddr, "ws-addr", "127.0.0.1:8088", "WebSocket服务器地址")
+	flag.BoolVar(&options.AIAnalysis, "ai", false, "AI分析")
+	flag.StringVar(&options.AIHunyuanToken, "hunyuan-token", "", "混元token")
+	flag.StringVar(&options.AIDeepSeekToken, "deepseek-token", "", "deepseek r1 token")
 	flag.Parse()
 	options.configureOutput()
 	ShowBanner()

pkg/hunyuan/hunyuan.go → pkg/openai/openai.go

@@ -1,5 +1,5 @@
 // Package hunyuan 大模型接口实现
-package hunyuan
+package openai
 
 import (
 	"context"
@@ -11,8 +11,22 @@ import (
 	"github.com/openai/openai-go/option"
 )
 
-// HunyuanAI Golang实现
+// HunyuanAI Hunyuan 实现
 func HunyuanAI(prompt string, key string) (string, error) {
+	baseUrl := "https://api.hunyuan.cloud.tencent.com/v1/"
+	model := "hunyuan-lite"
+	return OpenAI(prompt, key, baseUrl, model)
+}
+
+// DeepSeekR1API DeepSeek AI接入
+func DeepSeekR1API(prompt, key string) (string, error) {
+	baseUrl := "https://api.deepseek.ai/v1/"
+	model := "deepseek-reasoner"
+	return OpenAI(prompt, key, baseUrl, model)
+}
+
+// OpenAI API接入
+func OpenAI(prompt, key, baseUrl, model string) (string, error) {
 	// 设置默认值
 	if key == "" {
 		key = os.Getenv("OPENAI_API_KEY")
@@ -21,13 +35,11 @@ func HunyuanAI(prompt string, key string) (string, error) {
 	if key == "" {
 		return "", errors.New("OPENAI_API_KEY is empty")
 	}
-	baseUrl := os.Getenv("OPENAI_BASE_URL") // 默认值
 	if baseUrl == "" {
-		baseUrl = "https://api.hunyuan.cloud.tencent.com/v1/"
+		baseUrl = os.Getenv("OPENAI_BASE_URL")
 	}
-	model := os.Getenv("OPENAI_MODEL")
 	if model == "" {
-		model = "hunyuan-lite"
+		model = os.Getenv("OPENAI_MODEL")
 	}
 	client := openai.NewClient(option.WithBaseURL(baseUrl), option.WithAPIKey(key))
 	ctx := context.Background()

pkg/hunyuan/hunyuan_test.go → pkg/openai/openai_test.go

@@ -1,4 +1,4 @@
-package hunyuan
+package openai
 
 import (
 	"github.com/stretchr/testify/assert"