CalvinBackup/CyberStrikeAI
1
0
Fork 0
mirror of https://github.com/Ed1s0nZ/CyberStrikeAI.git synced 2026-07-14 07:57:26 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
公明
2026-07-10 16:46:26 +08:00
committed by GitHub
parent 25a76a8c97
commit 0ff8c58fbd
18 changed files with 1461 additions and 122 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/handler/c2.go
+88 -14
View File
@@ -17,6 +17,7 @@ import (
"cyberstrike-ai/internal/audit"
"cyberstrike-ai/internal/c2"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/security"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
@@ -59,7 +60,7 @@ func (h *C2Handler) SetManager(m *c2.Manager) {
// ListListeners 获取监听器列表
func (h *C2Handler) ListListeners(c *gin.Context) {
listeners, err := h.mgr().DB().ListC2Listeners()
listeners, err := h.mgr().DB().ListC2ListenersForAccess(c2AccessFromContext(c))
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
@@ -109,6 +110,11 @@ func (h *C2Handler) CreateListener(c *gin.Context) {
c.JSON(code, gin.H{"error": err.Error()})
return
}
if session, ok := security.CurrentSession(c); ok {
listener.OwnerUserID = session.UserID
_ = h.mgr().DB().SetResourceOwner("c2_listener", listener.ID, session.UserID)
_ = h.mgr().DB().AssignResourceToUser(session.UserID, "c2_listener", listener.ID)
}
implantToken := listener.ImplantToken
listener.EncryptionKey = ""
listener.ImplantToken = ""
@@ -282,7 +288,7 @@ func (h *C2Handler) ListSessions(c *gin.Context) {
filter.Suspicious = true
}
sessions, err := h.mgr().DB().ListC2Sessions(filter)
sessions, err := h.mgr().DB().ListC2SessionsForAccess(filter, c2AccessFromContext(c))
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
@@ -304,10 +310,10 @@ func (h *C2Handler) GetSession(c *gin.Context) {
}
// 获取最近任务
tasks, _ := h.mgr().DB().ListC2Tasks(database.ListC2TasksFilter{
tasks, _ := h.mgr().DB().ListC2TasksForAccess(database.ListC2TasksFilter{
SessionID: id,
Limit: 20,
})
}, c2AccessFromContext(c))
c.JSON(http.StatusOK, gin.H{
"session": session,
@@ -341,7 +347,7 @@ func (h *C2Handler) DeleteSessions(c *gin.Context) {
c.JSON(http.StatusBadRequest, gin.H{"error": "ids is required"})
return
}
n, err := h.mgr().DB().DeleteC2SessionsByIDs(req.IDs)
n, err := h.mgr().DB().DeleteC2SessionsByIDsForAccess(req.IDs, c2AccessFromContext(c))
if err != nil {
if errors.Is(err, database.ErrNoValidC2SessionIDs) {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
@@ -433,24 +439,25 @@ func (h *C2Handler) ListTasks(c *gin.Context) {
}
}
tasks, err := h.mgr().DB().ListC2Tasks(filter)
access := c2AccessFromContext(c)
tasks, err := h.mgr().DB().ListC2TasksForAccess(filter, access)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
// 仪表盘「待审任务」为全局 queued/pending 数量，与列表 session 过滤无关
pendingN, _ := h.mgr().DB().CountC2TasksQueuedOrPending("")
pendingN, _ := h.mgr().DB().CountC2TasksQueuedOrPendingForAccess("", access)
if !paginated {
c.JSON(http.StatusOK, gin.H{
"tasks": tasks,
"pending_queued_count": pendingN,
"tasks": tasks,
"pending_queued_count": pendingN,
})
return
}
total, err := h.mgr().DB().CountC2Tasks(filter)
total, err := h.mgr().DB().CountC2TasksForAccess(filter, access)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
@@ -477,7 +484,7 @@ func (h *C2Handler) DeleteTasks(c *gin.Context) {
c.JSON(http.StatusBadRequest, gin.H{"error": "ids is required"})
return
}
n, err := h.mgr().DB().DeleteC2TasksByIDs(req.IDs)
n, err := h.mgr().DB().DeleteC2TasksByIDsForAccess(req.IDs, c2AccessFromContext(c))
if err != nil {
if errors.Is(err, database.ErrNoValidC2TaskIDs) {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
@@ -522,6 +529,13 @@ func (h *C2Handler) CreateTask(c *gin.Context) {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
if strings.TrimSpace(req.SessionID) == "" {
req.SessionID = strings.TrimSpace(c.Param("id"))
}
if !h.c2ResourceAllowed(c, "c2_session", req.SessionID) {
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
return
}
input := c2.EnqueueTaskInput{
SessionID: req.SessionID,
@@ -621,6 +635,10 @@ func (h *C2Handler) PayloadOneliner(c *gin.Context) {
c.JSON(http.StatusNotFound, gin.H{"error": "listener not found"})
return
}
if !h.c2ResourceAllowed(c, "c2_listener", req.ListenerID) {
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
return
}
host := c2.ResolveBeaconDialHost(listener, strings.TrimSpace(req.Host), h.logger, listener.ID)
@@ -684,6 +702,10 @@ func (h *C2Handler) PayloadBuild(c *gin.Context) {
c.JSON(http.StatusNotFound, gin.H{"error": "listener not found"})
return
}
if !h.c2ResourceAllowed(c, "c2_listener", req.ListenerID) {
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
return
}
builder := c2.NewPayloadBuilder(h.mgr(), h.logger, "", "")
input := c2.PayloadBuilderInput{
@@ -779,7 +801,8 @@ func (h *C2Handler) ListEvents(c *gin.Context) {
}
}
events, err := h.mgr().DB().ListC2Events(filter)
access := c2AccessFromContext(c)
events, err := h.mgr().DB().ListC2EventsForAccess(filter, access)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
@@ -788,7 +811,7 @@ func (h *C2Handler) ListEvents(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"events": events})
return
}
total, err := h.mgr().DB().CountC2Events(filter)
total, err := h.mgr().DB().CountC2EventsForAccess(filter, access)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
@@ -814,7 +837,7 @@ func (h *C2Handler) DeleteEvents(c *gin.Context) {
c.JSON(http.StatusBadRequest, gin.H{"error": "ids is required"})
return
}
n, err := h.mgr().DB().DeleteC2EventsByIDs(req.IDs)
n, err := h.mgr().DB().DeleteC2EventsByIDsForAccess(req.IDs, c2AccessFromContext(c))
if err != nil {
if errors.Is(err, database.ErrNoValidC2EventIDs) {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
@@ -851,6 +874,9 @@ func (h *C2Handler) EventStream(c *gin.Context) {
if !ok {
return false
}
if !h.c2EventAllowed(c, e) {
return true
}
data, _ := json.Marshal(e)
fmt.Fprintf(w, "data: %s\n\n", data)
return true
@@ -964,6 +990,10 @@ func (h *C2Handler) UploadFileForImplant(c *gin.Context) {
c.JSON(http.StatusBadRequest, gin.H{"error": "session_id and remote_path required"})
return
}
if !h.c2ResourceAllowed(c, "c2_session", sessionID) {
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
return
}
file, header, err := c.Request.FormFile("file")
if err != nil {
@@ -1018,6 +1048,10 @@ func (h *C2Handler) ListFiles(c *gin.Context) {
c.JSON(http.StatusBadRequest, gin.H{"error": "session_id required"})
return
}
if !h.c2ResourceAllowed(c, "c2_session", sessionID) {
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
return
}
files, err := h.mgr().DB().ListC2FilesBySession(sessionID)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
@@ -1038,6 +1072,10 @@ func (h *C2Handler) DownloadResultFile(c *gin.Context) {
c.JSON(http.StatusNotFound, gin.H{"error": "task not found"})
return
}
if !h.c2ResourceAllowed(c, "c2_task", taskID) {
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
return
}
if task.ResultBlobPath == "" {
c.JSON(http.StatusNotFound, gin.H{"error": "no result file for this task"})
return
@@ -1053,6 +1091,42 @@ func osCreate(path string) (*os.File, error) {
return os.Create(path)
}
func c2AccessFromContext(c *gin.Context) database.RBACListAccess {
session, ok := security.CurrentSession(c)
if !ok {
return database.RBACListAccess{}
}
return database.RBACListAccess{UserID: session.UserID, Scope: session.Scope}
}
func (h *C2Handler) c2ResourceAllowed(c *gin.Context, resourceType, resourceID string) bool {
session, ok := security.CurrentSession(c)
if !ok {
return false
}
return h.mgr().DB().UserCanAccessResource(session.UserID, session.Scope, resourceType, resourceID)
}
func (h *C2Handler) c2EventAllowed(c *gin.Context, e *c2.Event) bool {
if e == nil {
return false
}
session, ok := security.CurrentSession(c)
if !ok {
return false
}
if session.Scope == database.RBACScopeAll {
return true
}
if strings.TrimSpace(e.SessionID) != "" {
return h.mgr().DB().UserCanAccessResource(session.UserID, session.Scope, "c2_session", e.SessionID)
}
if strings.TrimSpace(e.TaskID) != "" {
return h.mgr().DB().UserCanAccessResource(session.UserID, session.Scope, "c2_task", e.TaskID)
}
return false
}
// ============================================================================
// 辅助函数（firstNonEmpty 已在 vulnerability.go 中定义）
// ============================================================================