mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-05-15 21:08:01 +02:00
公明
110 lines
2.9 KiB
Go
110 lines
2.9 KiB
Go
package c2
|
||
|
||
import (
|
||
"context"
|
||
"time"
|
||
|
||
"cyberstrike-ai/internal/database"
|
||
|
||
"go.uber.org/zap"
|
||
)
|
||
|
||
// SessionWatchdog 会话心跳看门狗：周期扫描所有 active/sleeping 会话，
|
||
// 把超过 (sleep * (1 + jitter%) * graceFactor + minGrace) 仍未心跳的标为 dead。
|
||
//
|
||
// 设计要点：
|
||
// - 单 goroutine + ticker，避免对每个会话开 timer，session 数量大时也线性 OK；
|
||
// - 阈值随会话自身 sleep/jitter 自适应（sleep=300s 的会话不能用 sleep=5s 的判定）；
|
||
// - 全局最小宽限期 minGrace 避免 sleep 配置错误的会话被误判；
|
||
// - 不读 implant_uuid，纯按 last_check_in 字段，与 listener 类型解耦。
|
||
type SessionWatchdog struct {
|
||
manager *Manager
|
||
logger *zap.Logger
|
||
interval time.Duration // 扫描周期，默认 15s
|
||
minGrace time.Duration // 最小宽限期，默认 30s
|
||
gracePct float64 // 心跳超时倍数，默认 3.0（即 3 倍 sleep 周期没心跳算掉线）
|
||
stopCh chan struct{}
|
||
}
|
||
|
||
// NewSessionWatchdog 创建看门狗
|
||
func NewSessionWatchdog(m *Manager) *SessionWatchdog {
|
||
return &SessionWatchdog{
|
||
manager: m,
|
||
logger: m.Logger().With(zap.String("component", "c2-watchdog")),
|
||
interval: 15 * time.Second,
|
||
minGrace: 30 * time.Second,
|
||
gracePct: 3.0,
|
||
stopCh: make(chan struct{}),
|
||
}
|
||
}
|
||
|
||
// Run 阻塞执行，直到 ctx.Done() 或 Stop()
|
||
func (w *SessionWatchdog) Run(ctx context.Context) {
|
||
t := time.NewTicker(w.interval)
|
||
defer t.Stop()
|
||
for {
|
||
select {
|
||
case <-ctx.Done():
|
||
return
|
||
case <-w.stopCh:
|
||
return
|
||
case <-t.C:
|
||
w.tick()
|
||
}
|
||
}
|
||
}
|
||
|
||
// Stop 停止
|
||
func (w *SessionWatchdog) Stop() {
|
||
select {
|
||
case <-w.stopCh:
|
||
default:
|
||
close(w.stopCh)
|
||
}
|
||
}
|
||
|
||
func (w *SessionWatchdog) tick() {
|
||
now := time.Now()
|
||
for _, status := range []string{string(SessionActive), string(SessionSleeping)} {
|
||
sessions, err := w.manager.DB().ListC2Sessions(database.ListC2SessionsFilter{Status: status})
|
||
if err != nil {
|
||
w.logger.Warn("watchdog 列表查询失败", zap.Error(err))
|
||
continue
|
||
}
|
||
for _, s := range sessions {
|
||
if w.isStale(s, now) {
|
||
if err := w.manager.MarkSessionDead(s.ID); err != nil {
|
||
w.logger.Warn("标记会话掉线失败", zap.String("session_id", s.ID), zap.Error(err))
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// isStale 判断会话是否超时
|
||
func (w *SessionWatchdog) isStale(s *database.C2Session, now time.Time) bool {
|
||
// 无心跳记录：以 first_seen_at 兜底
|
||
last := s.LastCheckIn
|
||
if last.IsZero() {
|
||
last = s.FirstSeenAt
|
||
}
|
||
sleep := s.SleepSeconds
|
||
if sleep <= 0 {
|
||
// TCP reverse 模式 sleep=0 → 用最小宽限期判定
|
||
return now.Sub(last) > w.minGrace*2
|
||
}
|
||
jitter := s.JitterPercent
|
||
if jitter < 0 {
|
||
jitter = 0
|
||
}
|
||
if jitter > 100 {
|
||
jitter = 100
|
||
}
|
||
// 阈值 = sleep * (1 + jitter%) * gracePct，再加 minGrace 兜底
|
||
expected := time.Duration(float64(sleep)*(1+float64(jitter)/100.0)*w.gracePct) * time.Second
|
||
if expected < w.minGrace {
|
||
expected = w.minGrace
|
||
}
|
||
return now.Sub(last) > expected
|
||
}
|