mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-07-13 23:47:24 +02:00
Add files via upload
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
@@ -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()
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// CRUD:C2 任务
|
||||
// ----------------------------------------------------------------------------
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// CRUD:C2 文件
|
||||
// ----------------------------------------------------------------------------
|
||||
@@ -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
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// CRUD:C2 事件审计
|
||||
// ----------------------------------------------------------------------------
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// CRUD:C2 Malleable Profile
|
||||
// ----------------------------------------------------------------------------
|
||||
|
||||
@@ -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] {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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()
|
||||
|
||||
@@ -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
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user