mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-01 08:37:41 +02:00
Add files via upload
This commit is contained in:
@@ -69,7 +69,7 @@ func (h *MonitorHandler) SetAgentHandler(ah *AgentHandler) {
|
||||
h.agentHandler = ah
|
||||
}
|
||||
|
||||
const monitorPageTopTools = 3
|
||||
const monitorPageTopTools = 6
|
||||
|
||||
// MonitorStatsSummary 工具调用汇总
|
||||
type MonitorStatsSummary struct {
|
||||
|
||||
@@ -2,6 +2,7 @@ package handler
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
@@ -47,6 +48,12 @@ type workflowDryRunRequest struct {
|
||||
Inputs map[string]interface{} `json:"inputs,omitempty"`
|
||||
}
|
||||
|
||||
type workflowGenerateDraftRequest struct {
|
||||
Prompt string `json:"prompt"`
|
||||
Options workflowrunner.DraftOptions `json:"options"`
|
||||
AvailableTools []workflowrunner.DraftTool `json:"available_tools,omitempty"`
|
||||
}
|
||||
|
||||
func (h *WorkflowHandler) List(c *gin.Context) {
|
||||
includeDisabled := strings.EqualFold(c.Query("includeDisabled"), "true") || c.Query("include_disabled") == "1"
|
||||
items, err := h.db.ListWorkflowDefinitions(includeDisabled)
|
||||
@@ -126,6 +133,44 @@ func (h *WorkflowHandler) DryRun(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"result": result})
|
||||
}
|
||||
|
||||
func (h *WorkflowHandler) GenerateDraft(c *gin.Context) {
|
||||
var req workflowGenerateDraftRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()})
|
||||
return
|
||||
}
|
||||
draftReq := workflowrunner.DraftRequest{
|
||||
Prompt: req.Prompt,
|
||||
Options: req.Options,
|
||||
AvailableTools: req.AvailableTools,
|
||||
}
|
||||
var result *workflowrunner.DraftResult
|
||||
var llmErr error
|
||||
if h.cfg != nil {
|
||||
if llmCfg, _, ok := h.cfg.ResolveAIChannel(""); ok && strings.TrimSpace(llmCfg.APIKey) != "" && strings.TrimSpace(llmCfg.Model) != "" {
|
||||
result, llmErr = workflowrunner.GenerateDraftFromLLM(c.Request.Context(), draftReq, llmCfg, h.logger)
|
||||
} else {
|
||||
llmErr = errors.New("AI 通道未配置 api_key 或 model")
|
||||
}
|
||||
} else {
|
||||
llmErr = errors.New("工作流生成器未加载平台 AI 配置")
|
||||
}
|
||||
if llmErr != nil {
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "大模型生成失败: " + llmErr.Error()})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "workflow", "generate_draft", "自然语言生成工作流草稿", "", "", map[string]interface{}{
|
||||
"generator": result.Generator,
|
||||
"nodes": result.Stats["nodes"],
|
||||
"edges": result.Stats["edges"],
|
||||
"high_risk": result.Audit.HighRisk,
|
||||
"savable": result.Audit.Savable,
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"result": result})
|
||||
}
|
||||
|
||||
func (h *WorkflowHandler) Update(c *gin.Context) {
|
||||
h.save(c, c.Param("id"))
|
||||
}
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -76,3 +77,26 @@ func TestWorkflowPackageHandlerInspectionAndCreateImport(t *testing.T) {
|
||||
t.Fatalf("saved=%#v", saved)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWorkflowHandlerGenerateDraft(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
h := NewWorkflowHandler(nil, zap.NewNop())
|
||||
body := bytes.NewBufferString(`{"prompt":"对目标资产做端口扫描,如果发现高危端口就执行加固脚本,最后输出报告","options":{"include_objective":true},"available_tools":[{"key":"nmap","name":"nmap","enabled":true}]}`)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/api/workflows/generate-draft", body)
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
h.GenerateDraft(c)
|
||||
if w.Code != http.StatusBadGateway {
|
||||
t.Fatalf("status=%d body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
var resp struct {
|
||||
Error string `json:"error"`
|
||||
}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(resp.Error, "大模型生成失败") {
|
||||
t.Fatalf("unexpected error: %#v", resp.Error)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user