Add files via upload

This commit is contained in:
公明
2026-07-10 16:42:35 +08:00
committed by GitHub
parent 225d9ba0d9
commit a2744d6936
14 changed files with 2631 additions and 62 deletions
+2 -2
View File
@@ -78,8 +78,8 @@ func buildAuditLogsWhere(filter ListAuditLogsFilter) (string, []interface{}) {
}
if q := strings.TrimSpace(filter.Query); q != "" {
like := "%" + q + "%"
conditions = append(conditions, "(message LIKE ? OR resource_id LIKE ? OR action LIKE ? OR category LIKE ?)")
args = append(args, like, like, like, like)
conditions = append(conditions, "(message LIKE ? OR resource_id LIKE ? OR action LIKE ? OR category LIKE ? OR detail_json LIKE ?)")
args = append(args, like, like, like, like, like)
}
return strings.Join(conditions, " AND "), args
}
+49 -1
View File
@@ -137,7 +137,7 @@ func (db *DB) GetBatchQueue(queueID string) (*BatchTaskQueueRow, error) {
// GetAllBatchQueues 获取所有批量任务队列
func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) {
rows, err := db.Query(
"SELECT "+batchQueueSelectColumns+" FROM batch_task_queues ORDER BY created_at DESC",
"SELECT " + batchQueueSelectColumns + " FROM batch_task_queues ORDER BY created_at DESC",
)
if err != nil {
return nil, fmt.Errorf("查询批量任务队列列表失败: %w", err)
@@ -168,6 +168,10 @@ func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) {
// ListBatchQueues 列出批量任务队列(支持筛选和分页)
func (db *DB) ListBatchQueues(limit, offset int, status, keyword string) ([]*BatchTaskQueueRow, error) {
return db.ListBatchQueuesForAccess(limit, offset, status, keyword, "", "")
}
func (db *DB) ListBatchQueuesForAccess(limit, offset int, status, keyword, userID, scope string) ([]*BatchTaskQueueRow, error) {
query := "SELECT " + batchQueueSelectColumns + " FROM batch_task_queues WHERE 1=1"
args := []interface{}{}
@@ -182,6 +186,26 @@ func (db *DB) ListBatchQueues(limit, offset int, status, keyword string) ([]*Bat
query += " AND (id LIKE ? OR title LIKE ?)"
args = append(args, "%"+keyword+"%", "%"+keyword+"%")
}
userID = strings.TrimSpace(userID)
if userID != "" && scope != RBACScopeAll {
query += ` AND (
owner_user_id = ?
OR EXISTS (
SELECT 1 FROM rbac_resource_assignments ra
WHERE ra.user_id = ? AND ra.resource_type = 'batch_task' AND ra.resource_id = batch_task_queues.id
)
OR (
project_id IS NOT NULL AND project_id <> '' AND (
EXISTS (SELECT 1 FROM projects p WHERE p.id = batch_task_queues.project_id AND p.owner_user_id = ?)
OR EXISTS (
SELECT 1 FROM rbac_resource_assignments pra
WHERE pra.user_id = ? AND pra.resource_type = 'project' AND pra.resource_id = batch_task_queues.project_id
)
)
)
)`
args = append(args, userID, userID, userID, userID)
}
query += " ORDER BY created_at DESC LIMIT ? OFFSET ?"
args = append(args, limit, offset)
@@ -216,6 +240,10 @@ func (db *DB) ListBatchQueues(limit, offset int, status, keyword string) ([]*Bat
// CountBatchQueues 统计批量任务队列总数(支持筛选条件)
func (db *DB) CountBatchQueues(status, keyword string) (int, error) {
return db.CountBatchQueuesForAccess(status, keyword, "", "")
}
func (db *DB) CountBatchQueuesForAccess(status, keyword, userID, scope string) (int, error) {
query := "SELECT COUNT(*) FROM batch_task_queues WHERE 1=1"
args := []interface{}{}
@@ -230,6 +258,26 @@ func (db *DB) CountBatchQueues(status, keyword string) (int, error) {
query += " AND (id LIKE ? OR title LIKE ?)"
args = append(args, "%"+keyword+"%", "%"+keyword+"%")
}
userID = strings.TrimSpace(userID)
if userID != "" && scope != RBACScopeAll {
query += ` AND (
owner_user_id = ?
OR EXISTS (
SELECT 1 FROM rbac_resource_assignments ra
WHERE ra.user_id = ? AND ra.resource_type = 'batch_task' AND ra.resource_id = batch_task_queues.id
)
OR (
project_id IS NOT NULL AND project_id <> '' AND (
EXISTS (SELECT 1 FROM projects p WHERE p.id = batch_task_queues.project_id AND p.owner_user_id = ?)
OR EXISTS (
SELECT 1 FROM rbac_resource_assignments pra
WHERE pra.user_id = ? AND pra.resource_type = 'project' AND pra.resource_id = batch_task_queues.project_id
)
)
)
)`
args = append(args, userID, userID, userID, userID)
}
var count int
err := db.QueryRow(query, args...).Scan(&count)
+507 -32
View File
@@ -45,20 +45,21 @@ func validC2TextIDForDelete(id string) bool {
// C2Listener 监听器实体
type C2Listener struct {
ID string `json:"id"`
Name string `json:"name"`
Type string `json:"type"` // tcp_reverse|http_beacon|https_beacon|websocket|dns
BindHost string `json:"bindHost"` // 默认 127.0.0.1
BindPort int `json:"bindPort"` // 1-65535
ProfileID string `json:"profileId"` // 可空:关联 c2_profiles.id
EncryptionKey string `json:"-"` // base64(AES-256),前端不返回
ImplantToken string `json:"-"` // beacon 携带的鉴权 token,前端不返回
Status string `json:"status"` // stopped|running|error
ConfigJSON string `json:"configJson"` // TLS 证书路径 / URI 模式 / 上限并发 等
Remark string `json:"remark"`
CreatedAt time.Time `json:"createdAt"`
ID string `json:"id"`
Name string `json:"name"`
Type string `json:"type"` // tcp_reverse|http_beacon|https_beacon|websocket|dns
BindHost string `json:"bindHost"` // 默认 127.0.0.1
BindPort int `json:"bindPort"` // 1-65535
ProfileID string `json:"profileId"` // 可空:关联 c2_profiles.id
EncryptionKey string `json:"-"` // base64(AES-256),前端不返回
ImplantToken string `json:"-"` // beacon 携带的鉴权 token,前端不返回
Status string `json:"status"` // stopped|running|error
ConfigJSON string `json:"configJson"` // TLS 证书路径 / URI 模式 / 上限并发 等
Remark string `json:"remark"`
OwnerUserID string `json:"ownerUserId,omitempty"`
CreatedAt time.Time `json:"createdAt"`
StartedAt *time.Time `json:"startedAt,omitempty"`
LastError string `json:"lastError,omitempty"`
LastError string `json:"lastError,omitempty"`
}
// C2Session 已上线会话
@@ -132,17 +133,17 @@ type C2Event struct {
// C2Profile Malleable Profile
type C2Profile struct {
ID string `json:"id"`
Name string `json:"name"`
UserAgent string `json:"userAgent"`
URIs []string `json:"uris"`
RequestHeaders map[string]string `json:"requestHeaders,omitempty"`
ResponseHeaders map[string]string `json:"responseHeaders,omitempty"`
BodyTemplate string `json:"bodyTemplate"`
JitterMinMS int `json:"jitterMinMs"`
JitterMaxMS int `json:"jitterMaxMs"`
Extra map[string]interface{} `json:"extra,omitempty"`
CreatedAt time.Time `json:"createdAt"`
ID string `json:"id"`
Name string `json:"name"`
UserAgent string `json:"userAgent"`
URIs []string `json:"uris"`
RequestHeaders map[string]string `json:"requestHeaders,omitempty"`
ResponseHeaders map[string]string `json:"responseHeaders,omitempty"`
BodyTemplate string `json:"bodyTemplate"`
JitterMinMS int `json:"jitterMinMs"`
JitterMaxMS int `json:"jitterMaxMs"`
Extra map[string]interface{} `json:"extra,omitempty"`
CreatedAt time.Time `json:"createdAt"`
}
// ----------------------------------------------------------------------------
@@ -165,12 +166,12 @@ func (db *DB) CreateC2Listener(l *C2Listener) error {
}
query := `
INSERT INTO c2_listeners (id, name, type, bind_host, bind_port, profile_id, encryption_key,
implant_token, status, config_json, remark, created_at, started_at, last_error)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
implant_token, status, config_json, remark, owner_user_id, created_at, started_at, last_error)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
`
_, err := db.Exec(query,
l.ID, l.Name, l.Type, l.BindHost, l.BindPort, l.ProfileID, l.EncryptionKey,
l.ImplantToken, l.Status, l.ConfigJSON, l.Remark, l.CreatedAt, l.StartedAt, l.LastError,
l.ImplantToken, l.Status, l.ConfigJSON, l.Remark, l.OwnerUserID, l.CreatedAt, l.StartedAt, l.LastError,
)
if err != nil {
db.logger.Error("创建 C2 监听器失败", zap.Error(err), zap.String("id", l.ID))
@@ -190,12 +191,12 @@ func (db *DB) UpdateC2Listener(l *C2Listener) error {
query := `
UPDATE c2_listeners SET
name = ?, type = ?, bind_host = ?, bind_port = ?, profile_id = ?, encryption_key = ?,
implant_token = ?, status = ?, config_json = ?, remark = ?, started_at = ?, last_error = ?
implant_token = ?, status = ?, config_json = ?, remark = ?, owner_user_id = ?, started_at = ?, last_error = ?
WHERE id = ?
`
res, err := db.Exec(query,
l.Name, l.Type, l.BindHost, l.BindPort, l.ProfileID, l.EncryptionKey,
l.ImplantToken, l.Status, l.ConfigJSON, l.Remark, l.StartedAt, l.LastError, l.ID,
l.ImplantToken, l.Status, l.ConfigJSON, l.Remark, l.OwnerUserID, l.StartedAt, l.LastError, l.ID,
)
if err != nil {
db.logger.Error("更新 C2 监听器失败", zap.Error(err), zap.String("id", l.ID))
@@ -231,7 +232,7 @@ func (db *DB) GetC2Listener(id string) (*C2Listener, error) {
SELECT id, name, type, bind_host, bind_port, COALESCE(profile_id, ''),
COALESCE(encryption_key, ''), COALESCE(implant_token, ''), status,
COALESCE(config_json, '{}'), COALESCE(remark, ''),
created_at, started_at, COALESCE(last_error, '')
COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '')
FROM c2_listeners WHERE id = ?
`
var l C2Listener
@@ -240,7 +241,7 @@ func (db *DB) GetC2Listener(id string) (*C2Listener, error) {
&l.ID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID,
&l.EncryptionKey, &l.ImplantToken, &l.Status,
&l.ConfigJSON, &l.Remark,
&l.CreatedAt, &startedAt, &l.LastError,
&l.OwnerUserID, &l.CreatedAt, &startedAt, &l.LastError,
)
if err == sql.ErrNoRows {
return nil, nil
@@ -261,7 +262,7 @@ func (db *DB) ListC2Listeners() ([]*C2Listener, error) {
SELECT id, name, type, bind_host, bind_port, COALESCE(profile_id, ''),
COALESCE(encryption_key, ''), COALESCE(implant_token, ''), status,
COALESCE(config_json, '{}'), COALESCE(remark, ''),
created_at, started_at, COALESCE(last_error, '')
COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '')
FROM c2_listeners ORDER BY created_at DESC
`
rows, err := db.Query(query)
@@ -277,6 +278,47 @@ func (db *DB) ListC2Listeners() ([]*C2Listener, error) {
&l.ID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID,
&l.EncryptionKey, &l.ImplantToken, &l.Status,
&l.ConfigJSON, &l.Remark,
&l.OwnerUserID, &l.CreatedAt, &startedAt, &l.LastError,
); err != nil {
db.logger.Warn("扫描 c2_listeners 行失败", zap.Error(err))
continue
}
if startedAt.Valid {
t := startedAt.Time
l.StartedAt = &t
}
list = append(list, &l)
}
return list, rows.Err()
}
// ListC2ListenersForAccess lists listeners visible to the resolved RBAC scope.
func (db *DB) ListC2ListenersForAccess(access RBACListAccess) ([]*C2Listener, error) {
conditions := []string{"1=1"}
args := []interface{}{}
appendC2ListenerAccessFilter(&conditions, &args, access)
query := `
SELECT id, name, type, bind_host, bind_port, COALESCE(profile_id, ''),
COALESCE(encryption_key, ''), COALESCE(implant_token, ''), status,
COALESCE(config_json, '{}'), COALESCE(remark, ''),
COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '')
FROM c2_listeners
WHERE ` + strings.Join(conditions, " AND ") + `
ORDER BY created_at DESC
`
rows, err := db.Query(query, args...)
if err != nil {
return nil, err
}
defer rows.Close()
var list []*C2Listener
for rows.Next() {
var l C2Listener
var startedAt sql.NullTime
if err := rows.Scan(
&l.ID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID,
&l.EncryptionKey, &l.ImplantToken, &l.Status,
&l.ConfigJSON, &l.Remark, &l.OwnerUserID,
&l.CreatedAt, &startedAt, &l.LastError,
); err != nil {
db.logger.Warn("扫描 c2_listeners 行失败", zap.Error(err))
@@ -291,6 +333,26 @@ func (db *DB) ListC2Listeners() ([]*C2Listener, error) {
return list, rows.Err()
}
func appendC2ListenerAccessFilter(conditions *[]string, args *[]interface{}, access RBACListAccess) {
if access.Scope == RBACScopeAll {
return
}
if access.UserID == "" {
*conditions = append(*conditions, "1=0")
return
}
clauses := []string{"owner_user_id = ?"}
*args = append(*args, access.UserID)
if access.Scope == RBACScopeAssigned {
clauses = append(clauses, `EXISTS (
SELECT 1 FROM rbac_resource_assignments ra
WHERE ra.user_id = ? AND ra.resource_type = 'c2_listener' AND ra.resource_id = c2_listeners.id
)`)
*args = append(*args, access.UserID)
}
*conditions = append(*conditions, "("+strings.Join(clauses, " OR ")+")")
}
// DeleteC2Listener 级联删除(会话/任务/文件/事件随之消失)
func (db *DB) DeleteC2Listener(id string) error {
res, err := db.Exec(`DELETE FROM c2_listeners WHERE id = ?`, id)
@@ -550,6 +612,109 @@ func (db *DB) ListC2Sessions(filter ListC2SessionsFilter) ([]*C2Session, error)
return list, rows.Err()
}
// ListC2SessionsForAccess lists sessions whose parent listener is visible.
func (db *DB) ListC2SessionsForAccess(filter ListC2SessionsFilter, access RBACListAccess) ([]*C2Session, error) {
conditions, args := buildC2SessionsWhere(filter)
appendC2SessionAccessFilter(&conditions, &args, access)
query := `
SELECT id, listener_id, implant_uuid, COALESCE(hostname,''), COALESCE(username,''),
COALESCE(os,''), COALESCE(arch,''), COALESCE(pid, 0), COALESCE(process_name,''),
COALESCE(is_admin, 0), COALESCE(internal_ip,''), COALESCE(external_ip,''),
COALESCE(user_agent,''), COALESCE(sleep_seconds, 5), COALESCE(jitter_percent, 0),
status, first_seen_at, last_check_in, COALESCE(metadata_json, '{}'),
COALESCE(note, '')
FROM c2_sessions
WHERE ` + strings.Join(conditions, " AND ") + `
ORDER BY last_check_in DESC
`
if filter.Limit > 0 {
query += fmt.Sprintf(" LIMIT %d", filter.Limit)
}
rows, err := db.Query(query, args...)
if err != nil {
return nil, err
}
defer rows.Close()
return db.scanC2SessionRows(rows)
}
func buildC2SessionsWhere(filter ListC2SessionsFilter) ([]string, []interface{}) {
conditions := []string{"1=1"}
args := []interface{}{}
if filter.ListenerID != "" {
conditions = append(conditions, "listener_id = ?")
args = append(args, filter.ListenerID)
}
if filter.Status != "" {
conditions = append(conditions, "status = ?")
args = append(args, filter.Status)
}
if filter.OS != "" {
conditions = append(conditions, "os = ?")
args = append(args, filter.OS)
}
if filter.Search != "" {
conditions = append(conditions, "(hostname LIKE ? OR username LIKE ? OR internal_ip LIKE ?)")
kw := "%" + filter.Search + "%"
args = append(args, kw, kw, kw)
}
if filter.Suspicious {
conditions = append(conditions, `status = 'dead' AND (
hostname LIKE 'tcp_%' OR LOWER(COALESCE(username,'')) = 'unknown' OR COALESCE(pid, 0) = 0
)`)
}
return conditions, args
}
func (db *DB) scanC2SessionRows(rows *sql.Rows) ([]*C2Session, error) {
var list []*C2Session
for rows.Next() {
var s C2Session
var isAdminInt int
var metadataJSON string
if err := rows.Scan(
&s.ID, &s.ListenerID, &s.ImplantUUID, &s.Hostname, &s.Username,
&s.OS, &s.Arch, &s.PID, &s.ProcessName,
&isAdminInt, &s.InternalIP, &s.ExternalIP,
&s.UserAgent, &s.SleepSeconds, &s.JitterPercent,
&s.Status, &s.FirstSeenAt, &s.LastCheckIn, &metadataJSON,
&s.Note,
); err != nil {
db.logger.Warn("扫描 c2_sessions 行失败", zap.Error(err))
continue
}
s.IsAdmin = isAdminInt != 0
if metadataJSON != "" && metadataJSON != "{}" {
_ = json.Unmarshal([]byte(metadataJSON), &s.Metadata)
}
list = append(list, &s)
}
return list, rows.Err()
}
func appendC2SessionAccessFilter(conditions *[]string, args *[]interface{}, access RBACListAccess) {
if access.Scope == RBACScopeAll {
return
}
if access.UserID == "" {
*conditions = append(*conditions, "1=0")
return
}
clauses := []string{`EXISTS (
SELECT 1 FROM c2_listeners
WHERE c2_listeners.id = c2_sessions.listener_id AND c2_listeners.owner_user_id = ?
)`}
*args = append(*args, access.UserID)
if access.Scope == RBACScopeAssigned {
clauses = append(clauses, `EXISTS (
SELECT 1 FROM rbac_resource_assignments ra
WHERE ra.user_id = ? AND ra.resource_type = 'c2_listener' AND ra.resource_id = c2_sessions.listener_id
)`)
*args = append(*args, access.UserID)
}
*conditions = append(*conditions, "("+strings.Join(clauses, " OR ")+")")
}
// DeleteC2Session 级联删除其 tasks/files
func (db *DB) DeleteC2Session(id string) error {
res, err := db.Exec(`DELETE FROM c2_sessions WHERE id = ?`, id)
@@ -601,6 +766,29 @@ func (db *DB) DeleteC2SessionsByIDs(ids []string) (int64, error) {
return res.RowsAffected()
}
func (db *DB) DeleteC2SessionsByIDsForAccess(ids []string, access RBACListAccess) (int64, error) {
if access.Scope == RBACScopeAll {
return db.DeleteC2SessionsByIDs(ids)
}
clean := cleanC2IDs(ids)
if len(clean) == 0 {
return 0, ErrNoValidC2SessionIDs
}
placeholders := strings.Repeat("?,", len(clean)-1) + "?"
args := make([]interface{}, 0, len(clean)+2)
for _, id := range clean {
args = append(args, id)
}
conditions := []string{"id IN (" + placeholders + ")"}
appendC2SessionAccessFilter(&conditions, &args, access)
query := `DELETE FROM c2_sessions WHERE ` + strings.Join(conditions, " AND ")
res, err := db.Exec(query, args...)
if err != nil {
return 0, err
}
return res.RowsAffected()
}
// ----------------------------------------------------------------------------
// CRUDC2 任务
// ----------------------------------------------------------------------------
@@ -778,6 +966,39 @@ func buildC2TasksWhere(filter ListC2TasksFilter) (where string, args []interface
return strings.Join(conditions, " AND "), args
}
func appendC2TaskAccessFilter(conditions *[]string, args *[]interface{}, access RBACListAccess) {
if access.Scope == RBACScopeAll {
return
}
if access.UserID == "" {
*conditions = append(*conditions, "1=0")
return
}
clauses := []string{`EXISTS (
SELECT 1 FROM c2_sessions s
JOIN c2_listeners l ON l.id = s.listener_id
WHERE s.id = c2_tasks.session_id AND l.owner_user_id = ?
)`}
*args = append(*args, access.UserID)
if access.Scope == RBACScopeAssigned {
clauses = append(clauses, `EXISTS (
SELECT 1 FROM c2_sessions s
JOIN rbac_resource_assignments ra ON ra.resource_id = s.listener_id
WHERE s.id = c2_tasks.session_id
AND ra.user_id = ? AND ra.resource_type = 'c2_listener'
)`)
*args = append(*args, access.UserID)
}
*conditions = append(*conditions, "("+strings.Join(clauses, " OR ")+")")
}
func buildC2TasksWhereForAccess(filter ListC2TasksFilter, access RBACListAccess) (string, []interface{}) {
where, args := buildC2TasksWhere(filter)
conditions := []string{where}
appendC2TaskAccessFilter(&conditions, &args, access)
return strings.Join(conditions, " AND "), args
}
// CountC2Tasks 与 ListC2Tasks 相同过滤条件下的记录总数
func (db *DB) CountC2Tasks(filter ListC2TasksFilter) (int64, error) {
where, args := buildC2TasksWhere(filter)
@@ -787,6 +1008,14 @@ func (db *DB) CountC2Tasks(filter ListC2TasksFilter) (int64, error) {
return n, err
}
func (db *DB) CountC2TasksForAccess(filter ListC2TasksFilter, access RBACListAccess) (int64, error) {
where, args := buildC2TasksWhereForAccess(filter, access)
query := `SELECT COUNT(*) FROM c2_tasks WHERE ` + where
var n int64
err := db.QueryRow(query, args...).Scan(&n)
return n, err
}
// CountC2TasksQueuedOrPending 统计 queued/pending 状态任务数(仪表盘「待审任务」)
func (db *DB) CountC2TasksQueuedOrPending(sessionID string) (int64, error) {
conditions := []string{"status IN ('queued', 'pending')"}
@@ -801,6 +1030,15 @@ func (db *DB) CountC2TasksQueuedOrPending(sessionID string) (int64, error) {
return n, err
}
func (db *DB) CountC2TasksQueuedOrPendingForAccess(sessionID string, access RBACListAccess) (int64, error) {
filter := ListC2TasksFilter{SessionID: sessionID}
where, args := buildC2TasksWhereForAccess(filter, access)
query := `SELECT COUNT(*) FROM c2_tasks WHERE status IN ('queued', 'pending') AND ` + where
var n int64
err := db.QueryRow(query, args...).Scan(&n)
return n, err
}
// ListC2Tasks 任务列表,按创建时间倒序
func (db *DB) ListC2Tasks(filter ListC2TasksFilter) ([]*C2Task, error) {
where, args := buildC2TasksWhere(filter)
@@ -866,6 +1104,74 @@ func (db *DB) ListC2Tasks(filter ListC2TasksFilter) ([]*C2Task, error) {
return list, rows.Err()
}
func (db *DB) ListC2TasksForAccess(filter ListC2TasksFilter, access RBACListAccess) ([]*C2Task, error) {
where, args := buildC2TasksWhereForAccess(filter, access)
query := `
SELECT id, session_id, task_type, COALESCE(payload_json, '{}'),
status, COALESCE(result_text, ''), COALESCE(result_blob_path, ''),
COALESCE(error, ''), COALESCE(source, 'manual'),
COALESCE(conversation_id, ''), COALESCE(approval_status, ''),
created_at, sent_at, started_at, completed_at, COALESCE(duration_ms, 0)
FROM c2_tasks
WHERE ` + where + `
ORDER BY created_at DESC
`
limit := filter.Limit
offset := filter.Offset
if offset < 0 {
offset = 0
}
if limit > 0 {
if limit > 1000 {
limit = 1000
}
query += ` LIMIT ? OFFSET ?`
args = append(args, limit, offset)
}
rows, err := db.Query(query, args...)
if err != nil {
return nil, err
}
defer rows.Close()
return db.scanC2TaskRows(rows)
}
func (db *DB) scanC2TaskRows(rows *sql.Rows) ([]*C2Task, error) {
var list []*C2Task
for rows.Next() {
var t C2Task
var payloadJSON string
var sentAt, startedAt, completedAt sql.NullTime
if err := rows.Scan(
&t.ID, &t.SessionID, &t.TaskType, &payloadJSON,
&t.Status, &t.ResultText, &t.ResultBlobPath,
&t.Error, &t.Source,
&t.ConversationID, &t.ApprovalStatus,
&t.CreatedAt, &sentAt, &startedAt, &completedAt, &t.DurationMS,
); err != nil {
db.logger.Warn("扫描 c2_tasks 行失败", zap.Error(err))
continue
}
if payloadJSON != "" && payloadJSON != "{}" {
_ = json.Unmarshal([]byte(payloadJSON), &t.Payload)
}
if sentAt.Valid {
x := sentAt.Time
t.SentAt = &x
}
if startedAt.Valid {
x := startedAt.Time
t.StartedAt = &x
}
if completedAt.Valid {
x := completedAt.Time
t.CompletedAt = &x
}
list = append(list, &t)
}
return list, rows.Err()
}
// PopQueuedC2Tasks 取出某会话所有 queued/approved 任务(用于 beacon 拉取),原子置为 sent
func (db *DB) PopQueuedC2Tasks(sessionID string, limit int) ([]*C2Task, error) {
if limit <= 0 {
@@ -978,6 +1284,29 @@ func (db *DB) DeleteC2TasksByIDs(ids []string) (int64, error) {
return res.RowsAffected()
}
func (db *DB) DeleteC2TasksByIDsForAccess(ids []string, access RBACListAccess) (int64, error) {
if access.Scope == RBACScopeAll {
return db.DeleteC2TasksByIDs(ids)
}
clean := cleanC2IDs(ids)
if len(clean) == 0 {
return 0, ErrNoValidC2TaskIDs
}
placeholders := strings.Repeat("?,", len(clean)-1) + "?"
args := make([]interface{}, 0, len(clean)+2)
for _, id := range clean {
args = append(args, id)
}
conditions := []string{"id IN (" + placeholders + ")"}
appendC2TaskAccessFilter(&conditions, &args, access)
query := `DELETE FROM c2_tasks WHERE ` + strings.Join(conditions, " AND ")
res, err := db.Exec(query, args...)
if err != nil {
return 0, err
}
return res.RowsAffected()
}
// ----------------------------------------------------------------------------
// CRUDC2 文件
// ----------------------------------------------------------------------------
@@ -1024,6 +1353,27 @@ func (db *DB) ListC2FilesBySession(sessionID string) ([]*C2File, error) {
return list, rows.Err()
}
func cleanC2IDs(ids []string) []string {
const maxBatch = 500
if len(ids) > maxBatch {
ids = ids[:maxBatch]
}
clean := make([]string, 0, len(ids))
seen := make(map[string]struct{}, len(ids))
for _, id := range ids {
id = strings.TrimSpace(id)
if !validC2TextIDForDelete(id) {
continue
}
if _, ok := seen[id]; ok {
continue
}
seen[id] = struct{}{}
clean = append(clean, id)
}
return clean
}
// ----------------------------------------------------------------------------
// CRUDC2 事件审计
// ----------------------------------------------------------------------------
@@ -1093,6 +1443,56 @@ func buildC2EventsWhere(filter ListC2EventsFilter) (where string, args []interfa
return strings.Join(conditions, " AND "), args
}
func appendC2EventAccessFilter(conditions *[]string, args *[]interface{}, access RBACListAccess) {
if access.Scope == RBACScopeAll {
return
}
if access.UserID == "" {
*conditions = append(*conditions, "1=0")
return
}
clauses := []string{`EXISTS (
SELECT 1 FROM c2_sessions s
JOIN c2_listeners l ON l.id = s.listener_id
WHERE s.id = c2_events.session_id AND l.owner_user_id = ?
)`}
*args = append(*args, access.UserID)
if access.Scope == RBACScopeAssigned {
clauses = append(clauses, `EXISTS (
SELECT 1 FROM c2_sessions s
JOIN rbac_resource_assignments ra ON ra.resource_id = s.listener_id
WHERE s.id = c2_events.session_id
AND ra.user_id = ? AND ra.resource_type = 'c2_listener'
)`)
*args = append(*args, access.UserID)
}
clauses = append(clauses, `EXISTS (
SELECT 1 FROM c2_tasks t
JOIN c2_sessions s ON s.id = t.session_id
JOIN c2_listeners l ON l.id = s.listener_id
WHERE t.id = c2_events.task_id AND l.owner_user_id = ?
)`)
*args = append(*args, access.UserID)
if access.Scope == RBACScopeAssigned {
clauses = append(clauses, `EXISTS (
SELECT 1 FROM c2_tasks t
JOIN c2_sessions s ON s.id = t.session_id
JOIN rbac_resource_assignments ra ON ra.resource_id = s.listener_id
WHERE t.id = c2_events.task_id
AND ra.user_id = ? AND ra.resource_type = 'c2_listener'
)`)
*args = append(*args, access.UserID)
}
*conditions = append(*conditions, "("+strings.Join(clauses, " OR ")+")")
}
func buildC2EventsWhereForAccess(filter ListC2EventsFilter, access RBACListAccess) (string, []interface{}) {
where, args := buildC2EventsWhere(filter)
conditions := []string{where}
appendC2EventAccessFilter(&conditions, &args, access)
return strings.Join(conditions, " AND "), args
}
// CountC2Events 与 ListC2Events 相同过滤条件下的记录总数
func (db *DB) CountC2Events(filter ListC2EventsFilter) (int64, error) {
where, args := buildC2EventsWhere(filter)
@@ -1102,6 +1502,14 @@ func (db *DB) CountC2Events(filter ListC2EventsFilter) (int64, error) {
return n, err
}
func (db *DB) CountC2EventsForAccess(filter ListC2EventsFilter, access RBACListAccess) (int64, error) {
where, args := buildC2EventsWhereForAccess(filter, access)
query := `SELECT COUNT(*) FROM c2_events WHERE ` + where
var n int64
err := db.QueryRow(query, args...).Scan(&n)
return n, err
}
// ListC2Events 事件查询,按创建时间倒序
func (db *DB) ListC2Events(filter ListC2EventsFilter) ([]*C2Event, error) {
where, args := buildC2EventsWhere(filter)
@@ -1143,6 +1551,50 @@ func (db *DB) ListC2Events(filter ListC2EventsFilter) ([]*C2Event, error) {
return list, rows.Err()
}
func (db *DB) ListC2EventsForAccess(filter ListC2EventsFilter, access RBACListAccess) ([]*C2Event, error) {
where, args := buildC2EventsWhereForAccess(filter, access)
limit := filter.Limit
if limit <= 0 || limit > 1000 {
limit = 200
}
offset := filter.Offset
if offset < 0 {
offset = 0
}
query := `
SELECT id, level, category, COALESCE(session_id, ''), COALESCE(task_id, ''),
message, COALESCE(data_json, ''), created_at
FROM c2_events
WHERE ` + where + `
ORDER BY created_at DESC
LIMIT ? OFFSET ?
`
args = append(args, limit, offset)
rows, err := db.Query(query, args...)
if err != nil {
return nil, err
}
defer rows.Close()
return scanC2EventRows(rows)
}
func scanC2EventRows(rows *sql.Rows) ([]*C2Event, error) {
var list []*C2Event
for rows.Next() {
var e C2Event
var dataJSON string
if err := rows.Scan(&e.ID, &e.Level, &e.Category, &e.SessionID, &e.TaskID,
&e.Message, &dataJSON, &e.CreatedAt); err != nil {
continue
}
if dataJSON != "" {
_ = json.Unmarshal([]byte(dataJSON), &e.Data)
}
list = append(list, &e)
}
return list, rows.Err()
}
// DeleteC2EventsByIDs 按主键批量删除事件,返回实际删除行数
func (db *DB) DeleteC2EventsByIDs(ids []string) (int64, error) {
if len(ids) == 0 {
@@ -1181,6 +1633,29 @@ func (db *DB) DeleteC2EventsByIDs(ids []string) (int64, error) {
return res.RowsAffected()
}
func (db *DB) DeleteC2EventsByIDsForAccess(ids []string, access RBACListAccess) (int64, error) {
if access.Scope == RBACScopeAll {
return db.DeleteC2EventsByIDs(ids)
}
clean := cleanC2IDs(ids)
if len(clean) == 0 {
return 0, ErrNoValidC2EventIDs
}
placeholders := strings.Repeat("?,", len(clean)-1) + "?"
args := make([]interface{}, 0, len(clean)+4)
for _, id := range clean {
args = append(args, id)
}
conditions := []string{"id IN (" + placeholders + ")"}
appendC2EventAccessFilter(&conditions, &args, access)
query := `DELETE FROM c2_events WHERE ` + strings.Join(conditions, " AND ")
res, err := db.Exec(query, args...)
if err != nil {
return 0, err
}
return res.RowsAffected()
}
// ----------------------------------------------------------------------------
// CRUDC2 Malleable Profile
// ----------------------------------------------------------------------------
+206 -9
View File
@@ -384,6 +384,31 @@ func appendConversationProjectFilter(where string, args []interface{}, projectID
return where + fmt.Sprintf(" AND %s = ?", col), append(args, pid)
}
func appendConversationAccessFilter(where string, args []interface{}, userID, scope, alias string) (string, []interface{}) {
userID = strings.TrimSpace(userID)
if userID == "" || scope == RBACScopeAll {
return where, args
}
prefix := ""
if alias != "" {
prefix = alias + "."
}
where += fmt.Sprintf(` AND (%sowner_user_id = ? OR EXISTS (
SELECT 1 FROM rbac_resource_assignments ra
WHERE ra.user_id = ? AND ra.resource_type = 'conversation' AND ra.resource_id = %sid
) OR EXISTS (
SELECT 1 FROM projects p
WHERE p.id = %sproject_id AND (
p.owner_user_id = ? OR EXISTS (
SELECT 1 FROM rbac_resource_assignments pra
WHERE pra.user_id = ? AND pra.resource_type = 'project' AND pra.resource_id = p.id
)
)
))`, prefix, prefix, prefix)
args = append(args, userID, userID, userID, userID)
return where, args
}
// CountConversations 统计对话数量。
func (db *DB) CountConversations(search, projectID string) (int, error) {
var count int
@@ -410,6 +435,33 @@ func (db *DB) CountConversations(search, projectID string) (int, error) {
return count, nil
}
func (db *DB) CountConversationsForAccess(search, projectID, userID, scope string) (int, error) {
var count int
var err error
if search != "" {
searchPattern := "%" + search + "%"
where := ` WHERE (c.title LIKE ?
OR EXISTS (SELECT 1 FROM messages m WHERE m.conversation_id = c.id AND m.content LIKE ?))`
args := []interface{}{searchPattern, searchPattern}
where, args = appendConversationProjectFilter(where, args, projectID, "c")
where, args = appendConversationAccessFilter(where, args, userID, scope, "c")
err = db.QueryRow(`SELECT COUNT(*) FROM conversations c`+where, args...).Scan(&count)
} else {
where := ""
args := []interface{}{}
where, args = appendConversationProjectFilter(where, args, projectID, "")
where, args = appendConversationAccessFilter(where, args, userID, scope, "")
if where != "" {
where = " WHERE" + strings.TrimPrefix(where, " AND")
}
err = db.QueryRow(`SELECT COUNT(*) FROM conversations`+where, args...).Scan(&count)
}
if err != nil {
return 0, fmt.Errorf("统计对话失败: %w", err)
}
return count, nil
}
func conversationOrderClause(sortBy, tableAlias string) string {
col := "updated_at"
if strings.TrimSpace(strings.ToLower(sortBy)) == "created_at" {
@@ -503,6 +555,81 @@ func (db *DB) ListConversations(limit, offset int, search, sortBy, projectID str
return conversations, nil
}
func (db *DB) ListConversationsForAccess(limit, offset int, search, sortBy, projectID, userID, scope string) ([]*Conversation, error) {
if scope == RBACScopeAll || strings.TrimSpace(userID) == "" {
return db.ListConversations(limit, offset, search, sortBy, projectID)
}
var rows *sql.Rows
var err error
if search != "" {
searchPattern := "%" + search + "%"
orderClause := conversationOrderClause(sortBy, "c")
where := ` WHERE (c.title LIKE ?
OR EXISTS (SELECT 1 FROM messages m WHERE m.conversation_id = c.id AND m.content LIKE ?))`
args := []interface{}{searchPattern, searchPattern}
where, args = appendConversationProjectFilter(where, args, projectID, "c")
where, args = appendConversationAccessFilter(where, args, userID, scope, "c")
args = append(args, limit, offset)
rows, err = db.Query(
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id
FROM conversations c`+where+`
`+orderClause+`
LIMIT ? OFFSET ?`, args...)
} else {
orderClause := conversationOrderClause(sortBy, "")
where := ""
args := []interface{}{}
where, args = appendConversationProjectFilter(where, args, projectID, "")
where, args = appendConversationAccessFilter(where, args, userID, scope, "")
if where != "" {
where = " WHERE" + strings.TrimPrefix(where, " AND")
}
args = append(args, limit, offset)
rows, err = db.Query(
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id FROM conversations"+where+" "+orderClause+" LIMIT ? OFFSET ?",
args...)
}
if err != nil {
return nil, fmt.Errorf("查询对话列表失败: %w", err)
}
defer rows.Close()
return scanConversationRows(rows)
}
func scanConversationRows(rows *sql.Rows) ([]*Conversation, error) {
var conversations []*Conversation
for rows.Next() {
var conv Conversation
var createdAt, updatedAt string
var pinned int
var projectID sql.NullString
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID); err != nil {
return nil, fmt.Errorf("扫描对话失败: %w", err)
}
if projectID.Valid {
conv.ProjectID = strings.TrimSpace(projectID.String)
}
var err1, err2 error
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
if err1 != nil {
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05", createdAt)
}
if err1 != nil {
conv.CreatedAt, _ = time.Parse(time.RFC3339, createdAt)
}
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt)
if err2 != nil {
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05", updatedAt)
}
if err2 != nil {
conv.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt)
}
conv.Pinned = pinned != 0
conversations = append(conversations, &conv)
}
return conversations, rows.Err()
}
const ungroupedConversationsSQL = `
FROM conversations c
WHERE NOT EXISTS (
@@ -521,6 +648,18 @@ func (db *DB) CountUngroupedConversations(projectID string) (int, error) {
return count, nil
}
func (db *DB) CountUngroupedConversationsForAccess(projectID, userID, scope string) (int, error) {
where := ungroupedConversationsSQL
args := []interface{}{}
where, args = appendConversationProjectFilter(where, args, projectID, "c")
where, args = appendConversationAccessFilter(where, args, userID, scope, "c")
var count int
if err := db.QueryRow(`SELECT COUNT(*) `+where, args...).Scan(&count); err != nil {
return 0, fmt.Errorf("统计未分组对话失败: %w", err)
}
return count, nil
}
// ListUngroupedConversations 列出不在任何分组中的对话(最近对话侧栏)。
func (db *DB) ListUngroupedConversations(limit, offset int, sortBy, projectID string) ([]*Conversation, error) {
orderClause := conversationOrderClause(sortBy, "c")
@@ -578,6 +717,30 @@ func (db *DB) ListUngroupedConversations(limit, offset int, sortBy, projectID st
return conversations, rows.Err()
}
func (db *DB) ListUngroupedConversationsForAccess(limit, offset int, sortBy, projectID, userID, scope string) ([]*Conversation, error) {
if scope == RBACScopeAll || strings.TrimSpace(userID) == "" {
return db.ListUngroupedConversations(limit, offset, sortBy, projectID)
}
orderClause := conversationOrderClause(sortBy, "c")
where := ungroupedConversationsSQL
args := []interface{}{}
where, args = appendConversationProjectFilter(where, args, projectID, "c")
where, args = appendConversationAccessFilter(where, args, userID, scope, "c")
args = append(args, limit, offset)
rows, err := db.Query(
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id `+
where+`
`+orderClause+`
LIMIT ? OFFSET ?`,
args...,
)
if err != nil {
return nil, fmt.Errorf("查询未分组对话失败: %w", err)
}
defer rows.Close()
return scanConversationRows(rows)
}
// GetConversationTitle 获取对话标题(轻量查询,不加载消息)
func (db *DB) GetConversationTitle(id string) (string, error) {
var title string
@@ -1280,8 +1443,13 @@ func (db *DB) GetProcessDetailsSummary(messageID string) (*ProcessDetailsSummary
return nil, fmt.Errorf("查询工具执行摘要失败: %w", err)
}
seenExecIDs := make(map[string]bool)
toolIndexByCallID := make(map[string]int)
unmatchedToolIdx := 0
// A provider may reuse a fallback toolCallId across streaming rounds. Keep a
// FIFO per ID instead of a single index so every persisted call gets at most
// one result. Results without an ID fall back to the oldest unmatched call.
toolIndexesByCallID := make(map[string][]int)
lastMatchedToolIndexByCallID := make(map[string]int)
matchedToolIndexes := make([]bool, 0)
nextUnmatchedToolIdx := 0
for execRows.Next() {
var detailID string
var eventType string
@@ -1320,24 +1488,52 @@ func (db *DB) GetProcessDetailsSummary(messageID string) (*ProcessDetailsSummary
ProcessDetailID: strings.TrimSpace(detailID),
ToolName: toolName,
ToolCallID: toolCallID,
Status: "running",
// This summary is reconstructed from persisted history, not live
// execution state. Until a matching result is found the honest state
// is "result_missing", never "running".
Status: "result_missing",
})
matchedToolIndexes = append(matchedToolIndexes, false)
if toolCallID != "" {
toolIndexByCallID[toolCallID] = len(summary.ToolExecutions) - 1
toolIndexesByCallID[toolCallID] = append(toolIndexesByCallID[toolCallID], len(summary.ToolExecutions)-1)
}
}
if eventType == "tool_result" {
idx := -1
if toolCallID != "" {
if found, ok := toolIndexByCallID[toolCallID]; ok {
idx = found
queue := toolIndexesByCallID[toolCallID]
for len(queue) > 0 {
candidate := queue[0]
queue = queue[1:]
if candidate >= 0 && candidate < len(matchedToolIndexes) && !matchedToolIndexes[candidate] {
idx = candidate
break
}
}
toolIndexesByCallID[toolCallID] = queue
if idx < 0 {
// Multiple persisted result events for one call (for example an
// agent-facing reduced result replacing an earlier preview) update
// that call instead of consuming an unrelated FIFO entry.
if previous, ok := lastMatchedToolIndexByCallID[toolCallID]; ok {
idx = previous
}
}
}
if idx < 0 && unmatchedToolIdx < len(summary.ToolExecutions) {
idx = unmatchedToolIdx
unmatchedToolIdx++
if idx < 0 {
for nextUnmatchedToolIdx < len(matchedToolIndexes) && matchedToolIndexes[nextUnmatchedToolIdx] {
nextUnmatchedToolIdx++
}
if nextUnmatchedToolIdx < len(matchedToolIndexes) {
idx = nextUnmatchedToolIdx
nextUnmatchedToolIdx++
}
}
if idx >= 0 && idx < len(summary.ToolExecutions) {
matchedToolIndexes[idx] = true
if toolCallID != "" {
lastMatchedToolIndexByCallID[toolCallID] = idx
}
if summary.ToolExecutions[idx].ToolName == "" {
summary.ToolExecutions[idx].ToolName = toolName
}
@@ -1356,6 +1552,7 @@ func (db *DB) GetProcessDetailsSummary(messageID string) (*ProcessDetailsSummary
ExecutionID: execID,
Status: status,
})
matchedToolIndexes = append(matchedToolIndexes, true)
}
}
if execID != "" && !seenExecIDs[execID] {
+8
View File
@@ -478,6 +478,7 @@ func (db *DB) initTables() error {
status TEXT NOT NULL DEFAULT 'stopped',
config_json TEXT NOT NULL DEFAULT '{}',
remark TEXT NOT NULL DEFAULT '',
owner_user_id TEXT,
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
started_at DATETIME,
last_error TEXT
@@ -783,6 +784,10 @@ func (db *DB) initTables() error {
return fmt.Errorf("创建audit_logs表失败: %w", err)
}
if err := db.initRBACTables(); err != nil {
return fmt.Errorf("创建RBAC表失败: %w", err)
}
for tableName, ddl := range map[string]string{
"workflow_definitions": createWorkflowDefinitionsTable,
"workflow_runs": createWorkflowRunsTable,
@@ -853,6 +858,9 @@ func (db *DB) initTables() error {
if err := db.migrateWorkflowRunsTable(); err != nil {
db.logger.Warn("迁移workflow_runs表失败", zap.Error(err))
}
if err := db.migrateRBACOwnershipColumns(); err != nil {
db.logger.Warn("迁移RBAC资源归属字段失败", zap.Error(err))
}
if _, err := db.Exec(createIndexes); err != nil {
return fmt.Errorf("创建索引失败: %w", err)
@@ -0,0 +1,108 @@
package database
import (
"path/filepath"
"testing"
"go.uber.org/zap"
)
func TestProcessDetailsSummaryPairsMixedIdentifiedAndIDLessResults(t *testing.T) {
db, conversationID, messageID := setupProcessDetailsSummaryTest(t)
for _, id := range []string{"call-1", "call-2", "call-3", "call-4"} {
if err := db.AddProcessDetail(messageID, conversationID, "tool_call", "call", map[string]interface{}{
"toolName": "http-framework-test", "toolCallId": id,
}); err != nil {
t.Fatalf("AddProcessDetail(tool_call): %v", err)
}
}
results := []map[string]interface{}{
{"toolName": "http-framework-test", "toolCallId": "call-1", "success": true},
{"toolName": "http-framework-test", "toolCallId": "call-2", "success": true},
{"toolName": "http-framework-test", "success": true},
{"toolName": "http-framework-test", "success": true},
}
for _, result := range results {
if err := db.AddProcessDetail(messageID, conversationID, "tool_result", "result", result); err != nil {
t.Fatalf("AddProcessDetail(tool_result): %v", err)
}
}
summary, err := db.GetProcessDetailsSummary(messageID)
if err != nil {
t.Fatalf("GetProcessDetailsSummary: %v", err)
}
if len(summary.ToolExecutions) != 4 {
t.Fatalf("tool executions = %d, want 4", len(summary.ToolExecutions))
}
for i, execution := range summary.ToolExecutions {
if execution.Status != "completed" {
t.Fatalf("execution %d status = %q, want completed", i, execution.Status)
}
}
}
func TestProcessDetailsSummaryPairsRepeatedToolCallIDsFIFO(t *testing.T) {
db, conversationID, messageID := setupProcessDetailsSummaryTest(t)
for i := 0; i < 2; i++ {
if err := db.AddProcessDetail(messageID, conversationID, "tool_call", "call", map[string]interface{}{
"toolName": "execute", "toolCallId": "legacy-reused-id",
}); err != nil {
t.Fatalf("AddProcessDetail(tool_call): %v", err)
}
}
for i := 0; i < 2; i++ {
if err := db.AddProcessDetail(messageID, conversationID, "tool_result", "result", map[string]interface{}{
"toolName": "execute", "toolCallId": "legacy-reused-id", "success": true,
}); err != nil {
t.Fatalf("AddProcessDetail(tool_result): %v", err)
}
}
summary, err := db.GetProcessDetailsSummary(messageID)
if err != nil {
t.Fatalf("GetProcessDetailsSummary: %v", err)
}
if len(summary.ToolExecutions) != 2 {
t.Fatalf("tool executions = %d, want 2", len(summary.ToolExecutions))
}
for i, execution := range summary.ToolExecutions {
if execution.Status != "completed" {
t.Fatalf("execution %d status = %q, want completed", i, execution.Status)
}
}
}
func TestProcessDetailsSummaryDoesNotReportPersistedOrphanAsRunning(t *testing.T) {
db, conversationID, messageID := setupProcessDetailsSummaryTest(t)
if err := db.AddProcessDetail(messageID, conversationID, "tool_call", "call", map[string]interface{}{
"toolName": "execute", "toolCallId": "orphan",
}); err != nil {
t.Fatalf("AddProcessDetail(tool_call): %v", err)
}
summary, err := db.GetProcessDetailsSummary(messageID)
if err != nil {
t.Fatalf("GetProcessDetailsSummary: %v", err)
}
if len(summary.ToolExecutions) != 1 || summary.ToolExecutions[0].Status != "result_missing" {
t.Fatalf("tool executions = %#v, want result_missing", summary.ToolExecutions)
}
}
func setupProcessDetailsSummaryTest(t *testing.T) (*DB, string, string) {
t.Helper()
db, err := NewDB(filepath.Join(t.TempDir(), "process-details.db"), zap.NewNop())
if err != nil {
t.Fatalf("NewDB: %v", err)
}
t.Cleanup(func() { _ = db.Close() })
conversation, err := db.CreateConversation("process details", ConversationCreateMeta{})
if err != nil {
t.Fatalf("CreateConversation: %v", err)
}
message, err := db.AddMessage(conversation.ID, "assistant", "done", nil)
if err != nil {
t.Fatalf("AddMessage: %v", err)
}
return db, conversation.ID, message.ID
}
+66 -5
View File
@@ -58,11 +58,11 @@ type ProjectFact struct {
// ProjectFactListFilter 事实列表筛选。
type ProjectFactListFilter struct {
Category string
Confidence string
Search string
RelatedVulnerabilityID string
ExcludeDeprecated bool // 为 true 时排除 confidence=deprecated
Category string
Confidence string
Search string
RelatedVulnerabilityID string
ExcludeDeprecated bool // 为 true 时排除 confidence=deprecated
}
// CreateProject 创建项目。
@@ -143,6 +143,19 @@ func appendProjectListFilters(query string, args []interface{}, status, search s
return query, args
}
func appendProjectAccessFilter(query string, args []interface{}, userID, scope string) (string, []interface{}) {
userID = strings.TrimSpace(userID)
if userID == "" || scope == RBACScopeAll {
return query, args
}
query += ` AND (owner_user_id = ? OR EXISTS (
SELECT 1 FROM rbac_resource_assignments ra
WHERE ra.user_id = ? AND ra.resource_type = 'project' AND ra.resource_id = projects.id
))`
args = append(args, userID, userID)
return query, args
}
// CountProjects 统计项目数量。
func (db *DB) CountProjects(status, search string) (int, error) {
query := `SELECT COUNT(*) FROM projects WHERE 1=1`
@@ -155,6 +168,18 @@ func (db *DB) CountProjects(status, search string) (int, error) {
return count, nil
}
func (db *DB) CountProjectsForAccess(status, search, userID, scope string) (int, error) {
query := `SELECT COUNT(*) FROM projects WHERE 1=1`
args := []interface{}{}
query, args = appendProjectListFilters(query, args, status, search)
query, args = appendProjectAccessFilter(query, args, userID, scope)
var count int
if err := db.QueryRow(query, args...).Scan(&count); err != nil {
return 0, fmt.Errorf("统计项目失败: %w", err)
}
return count, nil
}
// ListProjects 列出项目。
func (db *DB) ListProjects(status, search string, limit, offset int) ([]*Project, error) {
if limit <= 0 {
@@ -189,6 +214,42 @@ func (db *DB) ListProjects(status, search string, limit, offset int) ([]*Project
return out, rows.Err()
}
func (db *DB) ListProjectsForAccess(status, search string, limit, offset int, userID, scope string) ([]*Project, error) {
if scope == RBACScopeAll || strings.TrimSpace(userID) == "" {
return db.ListProjects(status, search, limit, offset)
}
if limit <= 0 {
limit = 50
}
query := `SELECT id, name, COALESCE(description,''), COALESCE(scope_json,''), status, pinned, created_at, updated_at
FROM projects WHERE 1=1`
args := []interface{}{}
query, args = appendProjectListFilters(query, args, status, search)
query, args = appendProjectAccessFilter(query, args, userID, scope)
query += " ORDER BY pinned DESC, updated_at DESC LIMIT ? OFFSET ?"
args = append(args, limit, offset)
rows, err := db.Query(query, args...)
if err != nil {
return nil, fmt.Errorf("列出项目失败: %w", err)
}
defer rows.Close()
var out []*Project
for rows.Next() {
var p Project
var pinned int
var createdAt, updatedAt string
if err := rows.Scan(&p.ID, &p.Name, &p.Description, &p.ScopeJSON, &p.Status, &pinned, &createdAt, &updatedAt); err != nil {
return nil, err
}
p.Pinned = pinned != 0
p.CreatedAt = parseDBTime(createdAt)
p.UpdatedAt = parseDBTime(updatedAt)
out = append(out, &p)
}
return out, rows.Err()
}
// UpdateProject 更新项目。
func (db *DB) UpdateProject(p *Project) error {
p.UpdatedAt = time.Now()
+25 -4
View File
@@ -33,6 +33,10 @@ type ProjectDashboardSummary struct {
// GetProjectDashboardSummary 聚合跨项目近期事实(仅活跃项目、排除 deprecated)。
func (db *DB) GetProjectDashboardSummary(factLimit int) (*ProjectDashboardSummary, error) {
return db.GetProjectDashboardSummaryForAccess(factLimit, "", "")
}
func (db *DB) GetProjectDashboardSummaryForAccess(factLimit int, userID, scope string) (*ProjectDashboardSummary, error) {
if factLimit <= 0 {
factLimit = 5
}
@@ -44,25 +48,42 @@ func (db *DB) GetProjectDashboardSummary(factLimit int) (*ProjectDashboardSummar
RecentFacts: []ProjectDashboardFact{},
}
if err := db.QueryRow(`SELECT COUNT(*) FROM projects WHERE status = 'active'`).Scan(&out.Totals.ActiveProjects); err != nil {
projectAccess := ""
args := []interface{}{}
userID = strings.TrimSpace(userID)
if userID != "" && scope != RBACScopeAll {
projectAccess = ` AND (
p.owner_user_id = ?
OR EXISTS (
SELECT 1 FROM rbac_resource_assignments ra
WHERE ra.user_id = ? AND ra.resource_type = 'project' AND ra.resource_id = p.id
)
)`
args = append(args, userID, userID)
}
if err := db.QueryRow(`SELECT COUNT(*) FROM projects p WHERE p.status = 'active'`+projectAccess, args...).Scan(&out.Totals.ActiveProjects); err != nil {
return nil, fmt.Errorf("统计活跃项目失败: %w", err)
}
if err := db.QueryRow(
`SELECT COUNT(*) FROM project_facts f
INNER JOIN projects p ON p.id = f.project_id
WHERE f.confidence != 'deprecated' AND p.status = 'active'`,
WHERE f.confidence != 'deprecated' AND p.status = 'active'`+projectAccess,
args...,
).Scan(&out.Totals.TotalFacts); err != nil {
return nil, fmt.Errorf("统计事实失败: %w", err)
}
queryArgs := append([]interface{}{}, args...)
queryArgs = append(queryArgs, factLimit)
rows, err := db.Query(
`SELECT f.id, f.project_id, p.name, f.fact_key, f.category, f.summary, f.confidence, f.pinned, f.updated_at
FROM project_facts f
INNER JOIN projects p ON p.id = f.project_id
WHERE f.confidence != 'deprecated' AND p.status = 'active'
WHERE f.confidence != 'deprecated' AND p.status = 'active'`+projectAccess+`
ORDER BY f.pinned DESC, f.updated_at DESC
LIMIT ?`,
factLimit,
queryArgs...,
)
if err != nil {
return nil, fmt.Errorf("查询近期事实失败: %w", err)
File diff suppressed because it is too large Load Diff
+401
View File
@@ -0,0 +1,401 @@
package database
import (
"path/filepath"
"strings"
"testing"
"time"
"go.uber.org/zap"
)
func newRBACTestDB(t *testing.T) *DB {
t.Helper()
db, err := NewDB(filepath.Join(t.TempDir(), "rbac.db"), zap.NewNop())
if err != nil {
t.Fatalf("NewDB: %v", err)
}
t.Cleanup(func() { _ = db.Close() })
return db
}
func TestRBACProjectAndConversationListAccess(t *testing.T) {
db := newRBACTestDB(t)
p1, _ := db.CreateProject(&Project{Name: "visible"})
p2, _ := db.CreateProject(&Project{Name: "hidden"})
if err := db.SetResourceOwner("project", p1.ID, "u1"); err != nil {
t.Fatal(err)
}
c1, _ := db.CreateConversation("visible conv", ConversationCreateMeta{ProjectID: p1.ID})
c2, _ := db.CreateConversation("hidden conv", ConversationCreateMeta{ProjectID: p2.ID})
_ = db.SetResourceOwner("conversation", c1.ID, "u1")
_ = db.SetResourceOwner("conversation", c2.ID, "u2")
projects, err := db.ListProjectsForAccess("", "", 50, 0, "u1", RBACScopeOwn)
if err != nil {
t.Fatal(err)
}
if len(projects) != 1 || projects[0].ID != p1.ID {
t.Fatalf("projects = %#v, want only %s", projects, p1.ID)
}
convs, err := db.ListConversationsForAccess(50, 0, "", "", "", "u1", RBACScopeOwn)
if err != nil {
t.Fatal(err)
}
if len(convs) != 1 || convs[0].ID != c1.ID {
t.Fatalf("conversations = %#v, want only %s", convs, c1.ID)
}
}
func TestRBACVulnerabilityAccessInheritsProject(t *testing.T) {
db := newRBACTestDB(t)
user, err := db.CreateRBACUser("u1", "User 1", "hash", true, nil)
if err != nil {
t.Fatal(err)
}
p1, _ := db.CreateProject(&Project{Name: "visible"})
p2, _ := db.CreateProject(&Project{Name: "hidden"})
if err := db.AssignResourceToUser(user.ID, "project", p1.ID); err != nil {
t.Fatal(err)
}
v1, _ := db.CreateVulnerability(&Vulnerability{ProjectID: p1.ID, Title: "v1", Severity: "high"})
v2, _ := db.CreateVulnerability(&Vulnerability{ProjectID: p2.ID, Title: "v2", Severity: "high"})
items, err := db.ListVulnerabilitiesForAccess(50, 0, VulnerabilityListFilter{}, RBACListAccess{UserID: user.ID, Scope: RBACScopeAssigned})
if err != nil {
t.Fatal(err)
}
if len(items) != 1 || items[0].ID != v1.ID {
t.Fatalf("vulnerabilities = %#v, want only %s; hidden %s", items, v1.ID, v2.ID)
}
if !db.UserCanAccessResource(user.ID, RBACScopeAssigned, "vulnerability", v1.ID) {
t.Fatalf("expected project assignment to allow vulnerability detail")
}
if db.UserCanAccessResource(user.ID, RBACScopeAssigned, "vulnerability", v2.ID) {
t.Fatalf("unexpected access to hidden vulnerability")
}
}
func TestRBACConversationAccessInheritsProject(t *testing.T) {
db := newRBACTestDB(t)
user, err := db.CreateRBACUser("project-member", "Project Member", "hash", true, nil)
if err != nil {
t.Fatal(err)
}
project, err := db.CreateProject(&Project{Name: "assigned project"})
if err != nil {
t.Fatal(err)
}
conversation, err := db.CreateConversation("project conversation", ConversationCreateMeta{ProjectID: project.ID})
if err != nil {
t.Fatal(err)
}
if err := db.AssignResourceToUser(user.ID, "project", project.ID); err != nil {
t.Fatal(err)
}
rows, err := db.ListConversationsForAccess(50, 0, "", "", "", user.ID, RBACScopeAssigned)
if err != nil {
t.Fatal(err)
}
if len(rows) != 1 || rows[0].ID != conversation.ID {
t.Fatalf("conversations = %#v, want project conversation %s", rows, conversation.ID)
}
if !db.UserCanAccessResource(user.ID, RBACScopeAssigned, "conversation", conversation.ID) {
t.Fatal("expected project assignment to allow conversation detail")
}
}
func TestRBACBatchResourceAssignmentValidationAndAtomicity(t *testing.T) {
db := newRBACTestDB(t)
user, err := db.CreateRBACUser("batch-member", "Batch Member", "hash", true, nil)
if err != nil {
t.Fatal(err)
}
p1, err := db.CreateProject(&Project{Name: "p1"})
if err != nil {
t.Fatal(err)
}
p2, err := db.CreateProject(&Project{Name: "p2"})
if err != nil {
t.Fatal(err)
}
p3, err := db.CreateProject(&Project{Name: "p3"})
if err != nil {
t.Fatal(err)
}
options, err := db.ListAssignableRBACResources("project", "p1", 50)
if err != nil {
t.Fatal(err)
}
if len(options) != 1 || options[0].ID != p1.ID || options[0].Label != "p1" {
t.Fatalf("resource options = %#v, want p1", options)
}
firstPage, err := db.ListAssignableRBACResourcesPage("project", "", 2, 0)
if err != nil {
t.Fatal(err)
}
secondPage, err := db.ListAssignableRBACResourcesPage("project", "", 2, 2)
if err != nil {
t.Fatal(err)
}
if len(firstPage) != 2 || len(secondPage) != 1 {
t.Fatalf("paged resource options = %d + %d, want 2 + 1", len(firstPage), len(secondPage))
}
seen := map[string]bool{}
for _, option := range append(firstPage, secondPage...) {
seen[option.ID] = true
}
if !seen[p1.ID] || !seen[p2.ID] || !seen[p3.ID] {
t.Fatalf("paged resource options missed resources: %#v", seen)
}
if _, err := db.ListAssignableRBACResources("secret_table", "", 50); err == nil {
t.Fatal("expected unsupported picker resource type to fail")
}
if _, err := db.AssignResourcesToUser(user.ID, "unknown_type", []string{p1.ID}); err == nil {
t.Fatal("expected unsupported resource type to fail")
}
if _, err := db.AssignResourcesToUser(user.ID, "project", []string{p1.ID, "missing-project"}); err == nil {
t.Fatal("expected missing resource to fail the entire batch")
}
rows, err := db.ListRBACResourceAssignments(user.ID)
if err != nil {
t.Fatal(err)
}
if len(rows) != 0 {
t.Fatalf("partial grants persisted after failed batch: %#v", rows)
}
created, err := db.AssignResourcesToUser(user.ID, "project", []string{p1.ID, p1.ID, p2.ID})
if err != nil {
t.Fatal(err)
}
if created != 2 {
t.Fatalf("created = %d, want 2 unique grants", created)
}
created, err = db.AssignResourcesToUser(user.ID, "project", []string{p1.ID, p2.ID})
if err != nil {
t.Fatal(err)
}
if created != 0 {
t.Fatalf("idempotent retry created = %d, want 0", created)
}
rows, err = db.ListRBACResourceAssignments(user.ID)
if err != nil {
t.Fatal(err)
}
if len(rows) != 2 {
t.Fatalf("assignment count = %d, want 2", len(rows))
}
}
func TestRBACWebshellAndBatchListAccess(t *testing.T) {
db := newRBACTestDB(t)
ws1 := WebShellConnection{ID: "ws_visible", URL: "http://a", Type: "php", Method: "post", CreatedAt: time.Now()}
ws2 := WebShellConnection{ID: "ws_hidden", URL: "http://b", Type: "php", Method: "post", CreatedAt: time.Now()}
if err := db.CreateWebshellConnection(&ws1); err != nil {
t.Fatal(err)
}
if err := db.CreateWebshellConnection(&ws2); err != nil {
t.Fatal(err)
}
_ = db.SetResourceOwner("webshell", ws1.ID, "u1")
_ = db.SetResourceOwner("webshell", ws2.ID, "u2")
webshells, err := db.ListWebshellConnectionsForAccess("u1", RBACScopeOwn)
if err != nil {
t.Fatal(err)
}
if len(webshells) != 1 || webshells[0].ID != ws1.ID {
t.Fatalf("webshells = %#v, want only %s", webshells, ws1.ID)
}
if err := db.CreateBatchQueue("q_visible", "visible", "", "eino_single", "manual", "", nil, "", 1, []map[string]interface{}{{"id": "t1", "message": "a"}}); err != nil {
t.Fatal(err)
}
if err := db.CreateBatchQueue("q_hidden", "hidden", "", "eino_single", "manual", "", nil, "", 1, []map[string]interface{}{{"id": "t2", "message": "b"}}); err != nil {
t.Fatal(err)
}
_ = db.SetResourceOwner("batch_task", "q_visible", "u1")
_ = db.SetResourceOwner("batch_task", "q_hidden", "u2")
queues, err := db.ListBatchQueuesForAccess(50, 0, "all", "", "u1", RBACScopeOwn)
if err != nil {
t.Fatal(err)
}
if len(queues) != 1 || queues[0].ID != "q_visible" {
t.Fatalf("queues = %#v, want only q_visible", queues)
}
}
func TestRBACC2AccessInheritsListener(t *testing.T) {
db := newRBACTestDB(t)
now := time.Now()
l1 := &C2Listener{ID: "l_visible", Name: "visible", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9001, OwnerUserID: "u1", CreatedAt: now}
l2 := &C2Listener{ID: "l_hidden", Name: "hidden", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9002, OwnerUserID: "u2", CreatedAt: now}
if err := db.CreateC2Listener(l1); err != nil {
t.Fatal(err)
}
if err := db.CreateC2Listener(l2); err != nil {
t.Fatal(err)
}
if err := db.UpsertC2Session(&C2Session{ID: "s_visible", ListenerID: l1.ID, ImplantUUID: "implant-visible", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil {
t.Fatal(err)
}
if err := db.UpsertC2Session(&C2Session{ID: "s_hidden", ListenerID: l2.ID, ImplantUUID: "implant-hidden", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil {
t.Fatal(err)
}
if err := db.CreateC2Task(&C2Task{ID: "t_visible", SessionID: "s_visible", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
t.Fatal(err)
}
if err := db.CreateC2Task(&C2Task{ID: "t_hidden", SessionID: "s_hidden", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
t.Fatal(err)
}
if err := db.AppendC2Event(&C2Event{ID: "e_visible", Level: "info", Category: "task", SessionID: "s_visible", TaskID: "t_visible", Message: "visible", CreatedAt: now}); err != nil {
t.Fatal(err)
}
if err := db.AppendC2Event(&C2Event{ID: "e_hidden", Level: "info", Category: "task", SessionID: "s_hidden", TaskID: "t_hidden", Message: "hidden", CreatedAt: now}); err != nil {
t.Fatal(err)
}
access := RBACListAccess{UserID: "u1", Scope: RBACScopeOwn}
listeners, err := db.ListC2ListenersForAccess(access)
if err != nil {
t.Fatal(err)
}
if len(listeners) != 1 || listeners[0].ID != l1.ID {
t.Fatalf("listeners = %#v, want only %s", listeners, l1.ID)
}
sessions, err := db.ListC2SessionsForAccess(ListC2SessionsFilter{}, access)
if err != nil {
t.Fatal(err)
}
if len(sessions) != 1 || sessions[0].ID != "s_visible" {
t.Fatalf("sessions = %#v, want only s_visible", sessions)
}
tasks, err := db.ListC2TasksForAccess(ListC2TasksFilter{}, access)
if err != nil {
t.Fatal(err)
}
if len(tasks) != 1 || tasks[0].ID != "t_visible" {
t.Fatalf("tasks = %#v, want only t_visible", tasks)
}
events, err := db.ListC2EventsForAccess(ListC2EventsFilter{}, access)
if err != nil {
t.Fatal(err)
}
if len(events) != 1 || events[0].ID != "e_visible" {
t.Fatalf("events = %#v, want only e_visible", events)
}
if !db.UserCanAccessResource("u1", RBACScopeOwn, "c2_task", "t_visible") {
t.Fatalf("expected listener ownership to allow task detail")
}
if db.UserCanAccessResource("u1", RBACScopeOwn, "c2_task", "t_hidden") {
t.Fatalf("unexpected access to hidden task")
}
}
func TestRBACC2AssignedDeleteIsScoped(t *testing.T) {
db := newRBACTestDB(t)
user, err := db.CreateRBACUser("u1", "User 1", "hash", true, nil)
if err != nil {
t.Fatal(err)
}
now := time.Now()
if err := db.CreateC2Listener(&C2Listener{ID: "l_assigned", Name: "assigned", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9001, CreatedAt: now}); err != nil {
t.Fatal(err)
}
if err := db.CreateC2Listener(&C2Listener{ID: "l_hidden", Name: "hidden", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9002, CreatedAt: now}); err != nil {
t.Fatal(err)
}
if err := db.AssignResourceToUser(user.ID, "c2_listener", "l_assigned"); err != nil {
t.Fatal(err)
}
for _, row := range []struct {
sessionID string
listener string
taskID string
eventID string
}{
{"s_assigned", "l_assigned", "t_assigned", "e_assigned"},
{"s_hidden", "l_hidden", "t_hidden", "e_hidden"},
} {
if err := db.UpsertC2Session(&C2Session{ID: row.sessionID, ListenerID: row.listener, ImplantUUID: row.sessionID + "_uuid", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil {
t.Fatal(err)
}
if err := db.CreateC2Task(&C2Task{ID: row.taskID, SessionID: row.sessionID, TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
t.Fatal(err)
}
if err := db.AppendC2Event(&C2Event{ID: row.eventID, Level: "info", Category: "task", SessionID: row.sessionID, TaskID: row.taskID, Message: row.eventID, CreatedAt: now}); err != nil {
t.Fatal(err)
}
}
access := RBACListAccess{UserID: user.ID, Scope: RBACScopeAssigned}
n, err := db.DeleteC2TasksByIDsForAccess([]string{"t_assigned", "t_hidden"}, access)
if err != nil {
t.Fatal(err)
}
if n != 1 {
t.Fatalf("deleted tasks = %d, want 1", n)
}
if task, _ := db.GetC2Task("t_hidden"); task == nil {
t.Fatalf("hidden task was deleted")
}
n, err = db.DeleteC2EventsByIDsForAccess([]string{"e_assigned", "e_hidden"}, access)
if err != nil {
t.Fatal(err)
}
if n != 1 {
t.Fatalf("deleted events = %d, want 1", n)
}
hiddenEvents, err := db.ListC2Events(ListC2EventsFilter{TaskID: "t_hidden"})
if err != nil {
t.Fatal(err)
}
if len(hiddenEvents) != 1 {
t.Fatalf("hidden event count = %d, want 1", len(hiddenEvents))
}
}
func TestRBACAssignmentLabelsAndWeakTitles(t *testing.T) {
db := newRBACTestDB(t)
user, err := db.CreateRBACUser("label-member", "Label Member", "hash", true, nil)
if err != nil {
t.Fatal(err)
}
project, err := db.CreateProject(&Project{Name: "Alpha Project"})
if err != nil {
t.Fatal(err)
}
conversation, err := db.CreateConversation("1", ConversationCreateMeta{})
if err != nil {
t.Fatal(err)
}
if _, err := db.AssignResourcesToUser(user.ID, "project", []string{project.ID}); err != nil {
t.Fatal(err)
}
options, err := db.ListAssignableRBACResources("conversation", "", 10)
if err != nil {
t.Fatal(err)
}
if len(options) == 0 {
t.Fatal("expected conversation options")
}
for _, option := range options {
if option.ID == conversation.ID && !strings.Contains(option.Label, "1 ·") {
t.Fatalf("weak conversation label = %q, want suffix with short id", option.Label)
}
}
rows, err := db.ListRBACResourceAssignments(user.ID)
if err != nil {
t.Fatal(err)
}
if len(rows) != 1 {
t.Fatalf("assignments = %#v, want 1", rows)
}
if rows[0].ResourceLabel != "Alpha Project" {
t.Fatalf("assignment label = %q, want Alpha Project", rows[0].ResourceLabel)
}
}
+74 -7
View File
@@ -23,6 +23,11 @@ type VulnerabilityListFilter struct {
TaskTag string
}
type RBACListAccess struct {
UserID string
Scope string
}
func escapeVulnerabilityLikePattern(s string) string {
s = strings.ReplaceAll(s, `\`, `\\`)
s = strings.ReplaceAll(s, `%`, `\%`)
@@ -89,6 +94,40 @@ func (f VulnerabilityListFilter) appendWhere(query string, args []interface{}) (
return query, args
}
func appendVulnerabilityAccessFilter(query string, args []interface{}, access RBACListAccess) (string, []interface{}) {
userID := strings.TrimSpace(access.UserID)
if userID == "" || access.Scope == RBACScopeAll {
return query, args
}
query += ` AND (
owner_user_id = ?
OR EXISTS (
SELECT 1 FROM rbac_resource_assignments ra
WHERE ra.user_id = ? AND ra.resource_type = 'vulnerability' AND ra.resource_id = vulnerabilities.id
)
OR (
project_id IS NOT NULL AND project_id <> '' AND (
EXISTS (SELECT 1 FROM projects p WHERE p.id = vulnerabilities.project_id AND p.owner_user_id = ?)
OR EXISTS (
SELECT 1 FROM rbac_resource_assignments pra
WHERE pra.user_id = ? AND pra.resource_type = 'project' AND pra.resource_id = vulnerabilities.project_id
)
)
)
OR (
conversation_id IS NOT NULL AND conversation_id <> '' AND (
EXISTS (SELECT 1 FROM conversations c WHERE c.id = vulnerabilities.conversation_id AND c.owner_user_id = ?)
OR EXISTS (
SELECT 1 FROM rbac_resource_assignments cra
WHERE cra.user_id = ? AND cra.resource_type = 'conversation' AND cra.resource_id = vulnerabilities.conversation_id
)
)
)
)`
args = append(args, userID, userID, userID, userID, userID, userID)
return query, args
}
// Vulnerability 漏洞
type Vulnerability struct {
ID string `json:"id"`
@@ -190,6 +229,10 @@ func (db *DB) GetVulnerability(id string) (*Vulnerability, error) {
// ListVulnerabilities 列出漏洞
func (db *DB) ListVulnerabilities(limit, offset int, filter VulnerabilityListFilter) ([]*Vulnerability, error) {
return db.ListVulnerabilitiesForAccess(limit, offset, filter, RBACListAccess{})
}
func (db *DB) ListVulnerabilitiesForAccess(limit, offset int, filter VulnerabilityListFilter, access RBACListAccess) ([]*Vulnerability, error) {
query := `
SELECT id, COALESCE(conversation_id,''), COALESCE(project_id,''), title, description, severity, status, conversation_tag, task_tag,
vulnerability_type, target,
@@ -203,6 +246,7 @@ func (db *DB) ListVulnerabilities(limit, offset int, filter VulnerabilityListFil
`
args := []interface{}{}
query, args = filter.appendWhere(query, args)
query, args = appendVulnerabilityAccessFilter(query, args, access)
query += " ORDER BY created_at DESC LIMIT ? OFFSET ?"
args = append(args, limit, offset)
@@ -235,9 +279,14 @@ func (db *DB) ListVulnerabilities(limit, offset int, filter VulnerabilityListFil
// CountVulnerabilities 统计漏洞总数(支持筛选条件)
func (db *DB) CountVulnerabilities(filter VulnerabilityListFilter) (int, error) {
return db.CountVulnerabilitiesForAccess(filter, RBACListAccess{})
}
func (db *DB) CountVulnerabilitiesForAccess(filter VulnerabilityListFilter, access RBACListAccess) (int, error) {
query := "SELECT COUNT(*) FROM vulnerabilities WHERE 1=1"
args := []interface{}{}
query, args = filter.appendWhere(query, args)
query, args = appendVulnerabilityAccessFilter(query, args, access)
var count int
err := db.QueryRow(query, args...).Scan(&count)
@@ -275,6 +324,10 @@ func (db *DB) UpdateVulnerability(id string, vuln *Vulnerability) error {
// DeleteVulnerabilitiesByFilter 按筛选条件批量删除漏洞,返回实际删除条数
func (db *DB) DeleteVulnerabilitiesByFilter(filter VulnerabilityListFilter) (int64, error) {
return db.DeleteVulnerabilitiesByFilterForAccess(filter, RBACListAccess{})
}
func (db *DB) DeleteVulnerabilitiesByFilterForAccess(filter VulnerabilityListFilter, access RBACListAccess) (int64, error) {
tx, err := db.Begin()
if err != nil {
return 0, fmt.Errorf("开启事务失败: %w", err)
@@ -284,6 +337,7 @@ func (db *DB) DeleteVulnerabilitiesByFilter(filter VulnerabilityListFilter) (int
where := "WHERE 1=1"
args := []interface{}{}
where, args = filter.appendWhere(where, args)
where, args = appendVulnerabilityAccessFilter(where, args, access)
clearQuery := `UPDATE project_facts SET related_vulnerability_id = NULL
WHERE related_vulnerability_id IN (SELECT id FROM vulnerabilities ` + where + `)`
@@ -329,11 +383,16 @@ func (db *DB) DeleteVulnerability(id string) error {
// GetVulnerabilityStats 获取漏洞统计(筛选条件与 ListVulnerabilities / CountVulnerabilities 一致)
func (db *DB) GetVulnerabilityStats(filter VulnerabilityListFilter) (map[string]interface{}, error) {
return db.GetVulnerabilityStatsForAccess(filter, RBACListAccess{})
}
func (db *DB) GetVulnerabilityStatsForAccess(filter VulnerabilityListFilter, access RBACListAccess) (map[string]interface{}, error) {
stats := make(map[string]interface{})
where := "WHERE 1=1"
args := []interface{}{}
where, args = filter.appendWhere(where, args)
where, args = appendVulnerabilityAccessFilter(where, args, access)
// 总漏洞数
var totalCount int
@@ -389,6 +448,10 @@ func (db *DB) GetVulnerabilityStats(filter VulnerabilityListFilter) (map[string]
// GetVulnerabilityFilterOptions 获取漏洞筛选建议项
func (db *DB) GetVulnerabilityFilterOptions() (map[string][]string, error) {
return db.GetVulnerabilityFilterOptionsForAccess(RBACListAccess{})
}
func (db *DB) GetVulnerabilityFilterOptionsForAccess(access RBACListAccess) (map[string][]string, error) {
collect := func(query string, args ...interface{}) ([]string, error) {
rows, err := db.Query(query, args...)
if err != nil {
@@ -409,31 +472,35 @@ func (db *DB) GetVulnerabilityFilterOptions() (map[string][]string, error) {
return items, nil
}
vulnIDs, err := collect(`SELECT DISTINCT id FROM vulnerabilities ORDER BY created_at DESC LIMIT 500`)
where := "WHERE 1=1"
accessArgs := []interface{}{}
where, accessArgs = appendVulnerabilityAccessFilter(where, accessArgs, access)
vulnIDs, err := collect(`SELECT DISTINCT id FROM vulnerabilities `+where+` ORDER BY created_at DESC LIMIT 500`, accessArgs...)
if err != nil {
return nil, fmt.Errorf("查询漏洞ID建议失败: %w", err)
}
conversationIDs, err := collect(`SELECT DISTINCT conversation_id FROM vulnerabilities WHERE conversation_id IS NOT NULL AND conversation_id <> '' ORDER BY created_at DESC LIMIT 500`)
conversationIDs, err := collect(`SELECT DISTINCT conversation_id FROM vulnerabilities `+where+` AND conversation_id IS NOT NULL AND conversation_id <> '' ORDER BY created_at DESC LIMIT 500`, accessArgs...)
if err != nil {
return nil, fmt.Errorf("查询会话ID建议失败: %w", err)
}
taskIDs, err := collect(`SELECT DISTINCT id FROM batch_tasks WHERE id <> '' ORDER BY rowid DESC LIMIT 500`)
taskIDs, err := collect(`SELECT DISTINCT bt.id FROM batch_tasks bt JOIN vulnerabilities ON bt.conversation_id = vulnerabilities.conversation_id `+where+` AND bt.id <> '' ORDER BY bt.rowid DESC LIMIT 500`, accessArgs...)
if err != nil {
return nil, fmt.Errorf("查询任务ID建议失败: %w", err)
}
queueIDs, err := collect(`SELECT DISTINCT queue_id FROM batch_tasks WHERE queue_id <> '' ORDER BY rowid DESC LIMIT 500`)
queueIDs, err := collect(`SELECT DISTINCT bt.queue_id FROM batch_tasks bt JOIN vulnerabilities ON bt.conversation_id = vulnerabilities.conversation_id `+where+` AND bt.queue_id <> '' ORDER BY bt.rowid DESC LIMIT 500`, accessArgs...)
if err != nil {
return nil, fmt.Errorf("查询队列ID建议失败: %w", err)
}
conversationTags, err := collect(`SELECT DISTINCT conversation_tag FROM vulnerabilities WHERE conversation_tag IS NOT NULL AND conversation_tag <> '' ORDER BY conversation_tag LIMIT 500`)
conversationTags, err := collect(`SELECT DISTINCT conversation_tag FROM vulnerabilities `+where+` AND conversation_tag IS NOT NULL AND conversation_tag <> '' ORDER BY conversation_tag LIMIT 500`, accessArgs...)
if err != nil {
return nil, fmt.Errorf("查询对话标签建议失败: %w", err)
}
taskTags, err := collect(`SELECT DISTINCT task_tag FROM vulnerabilities WHERE task_tag IS NOT NULL AND task_tag <> '' ORDER BY task_tag LIMIT 500`)
taskTags, err := collect(`SELECT DISTINCT task_tag FROM vulnerabilities `+where+` AND task_tag IS NOT NULL AND task_tag <> '' ORDER BY task_tag LIMIT 500`, accessArgs...)
if err != nil {
return nil, fmt.Errorf("查询任务标签建议失败: %w", err)
}
projectIDs, err := collect(`SELECT DISTINCT project_id FROM vulnerabilities WHERE project_id IS NOT NULL AND project_id <> '' ORDER BY created_at DESC LIMIT 200`)
projectIDs, err := collect(`SELECT DISTINCT project_id FROM vulnerabilities `+where+` AND project_id IS NOT NULL AND project_id <> '' ORDER BY created_at DESC LIMIT 200`, accessArgs...)
if err != nil {
return nil, fmt.Errorf("查询项目ID建议失败: %w", err)
}
+20 -2
View File
@@ -2,6 +2,7 @@ package database
import (
"database/sql"
"strings"
"time"
"go.uber.org/zap"
@@ -59,13 +60,30 @@ func (db *DB) UpsertWebshellConnectionState(connectionID, stateJSON string) erro
// ListWebshellConnections 列出所有 WebShell 连接,按创建时间倒序
func (db *DB) ListWebshellConnections() ([]WebShellConnection, error) {
return db.ListWebshellConnectionsForAccess("", "")
}
func (db *DB) ListWebshellConnectionsForAccess(userID, scope string) ([]WebShellConnection, error) {
query := `
SELECT id, url, password, type, method, cmd_param, remark,
COALESCE(encoding, '') AS encoding, COALESCE(os, '') AS os, created_at
FROM webshell_connections
ORDER BY created_at DESC
WHERE 1=1
`
rows, err := db.Query(query)
args := []interface{}{}
userID = strings.TrimSpace(userID)
if userID != "" && scope != RBACScopeAll {
query += ` AND (
owner_user_id = ?
OR EXISTS (
SELECT 1 FROM rbac_resource_assignments ra
WHERE ra.user_id = ? AND ra.resource_type = 'webshell' AND ra.resource_id = webshell_connections.id
)
)`
args = append(args, userID, userID)
}
query += ` ORDER BY created_at DESC`
rows, err := db.Query(query, args...)
if err != nil {
db.logger.Error("查询 WebShell 连接列表失败", zap.Error(err))
return nil, err
+31
View File
@@ -4,8 +4,10 @@ import (
"context"
"fmt"
"strings"
"unicode/utf8"
"cyberstrike-ai/internal/agent"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/multiagent"
)
@@ -85,6 +87,11 @@ func runToolNode(ctx context.Context, args RunArgs, node graphNode, state *Workf
executionID = result.ExecutionID
isError = result.IsError
}
maxToolOutputBytes := config.MultiAgentEinoMiddlewareConfig{}.ReductionMaxLengthForTruncEffective()
if args.AppCfg != nil {
maxToolOutputBytes = args.AppCfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective()
}
output = truncateWorkflowToolOutput(output, maxToolOutputBytes, executionID)
out := toolOutputMap(node, output, toolName, toolArgs, executionID, isError)
if key := cfgString(node.Config, "output_key"); key != "" {
state.Outputs[key] = output
@@ -99,6 +106,30 @@ func runToolNode(ctx context.Context, args RunArgs, node graphNode, state *Workf
return out, true, "completed", ""
}
func truncateWorkflowToolOutput(output string, maxBytes int, executionID string) string {
if maxBytes <= 0 || len(output) <= maxBytes {
return output
}
marker := fmt.Sprintf("\n\n...[workflow tool output truncated; full result is stored in execution %s]...\n\n", strings.TrimSpace(executionID))
if strings.TrimSpace(executionID) == "" {
marker = "\n\n...[workflow tool output truncated; full result remains in the tool execution record]...\n\n"
}
budget := maxBytes - len(marker)
if budget <= 0 {
return marker
}
head := budget / 2
tail := budget - head
for head > 0 && !utf8.RuneStart(output[head]) {
head--
}
tailStart := len(output) - tail
for tailStart < len(output) && !utf8.RuneStart(output[tailStart]) {
tailStart++
}
return output[:head] + marker + output[tailStart:]
}
func runAgentNode(ctx context.Context, args RunArgs, node graphNode, state *WorkflowLocalState) (map[string]any, bool, string, string) {
if args.AppCfg == nil || args.Agent == nil {
errText := "Agent 节点执行失败:应用配置或 Agent 为空"
@@ -0,0 +1,23 @@
package workflow
import (
"strings"
"testing"
)
func TestTruncateWorkflowToolOutputBoundsBytesAndKeepsExecutionReference(t *testing.T) {
out := truncateWorkflowToolOutput(strings.Repeat("响应正文", 1000), 256, "exec-123")
if len(out) > 256 {
t.Fatalf("workflow output bytes=%d, want <=256", len(out))
}
if !strings.Contains(out, "exec-123") || !strings.Contains(out, "truncated") {
t.Fatalf("missing truncation reference: %q", out)
}
}
func TestTruncateWorkflowToolOutputLeavesBoundedContentUntouched(t *testing.T) {
const want = "small-result"
if got := truncateWorkflowToolOutput(want, 256, "exec-123"); got != want {
t.Fatalf("got %q want %q", got, want)
}
}