package logging import ( "fmt" "io" "log" "path/filepath" "runtime" "strings" "sync" "time" "github.com/fatih/color" ) // Logger 日志管理器 type Logger struct { mu sync.RWMutex config *LoggerConfig formatter LogFormatter handlers []LogHandler scanStatus *ScanStatus progressBar ProgressDisplay outputMutex *sync.Mutex initialized bool } // NewLogger 创建新的日志管理器 func NewLogger(config *LoggerConfig) *Logger { if config == nil { config = DefaultLoggerConfig() } logger := &Logger{ config: config, formatter: NewStandardFormatter(), handlers: make([]LogHandler, 0), scanStatus: NewScanStatus(), outputMutex: &sync.Mutex{}, initialized: true, } // 设置格式化器的开始时间 logger.formatter.SetStartTime(config.StartTime) // 添加默认的控制台处理器 consoleHandler := NewConsoleHandler(config) logger.AddHandler(consoleHandler) return logger } // SetFormatter 设置格式化器 func (l *Logger) SetFormatter(formatter LogFormatter) { l.mu.Lock() defer l.mu.Unlock() l.formatter = formatter l.formatter.SetStartTime(l.config.StartTime) } // AddHandler 添加日志处理器 func (l *Logger) AddHandler(handler LogHandler) { l.mu.Lock() defer l.mu.Unlock() l.handlers = append(l.handlers, handler) } // SetProgressBar 设置进度条显示 func (l *Logger) SetProgressBar(progressBar ProgressDisplay) { l.mu.Lock() defer l.mu.Unlock() l.progressBar = progressBar } // SetOutputMutex 设置输出互斥锁 func (l *Logger) SetOutputMutex(mutex *sync.Mutex) { l.mu.Lock() defer l.mu.Unlock() l.outputMutex = mutex } // Log 记录日志 func (l *Logger) Log(level LogLevel, content string, metadata ...map[string]interface{}) { if !l.shouldLog(level) { return } entry := &LogEntry{ Level: level, Time: time.Now(), Content: content, } // 添加元数据 if len(metadata) > 0 { entry.Metadata = metadata[0] } // 对于错误级别,自动添加调用者信息 if level == LevelError { if _, file, line, ok := runtime.Caller(2); ok { entry.Source = fmt.Sprintf("%s:%d", filepath.Base(file), line) entry.Content = fmt.Sprintf("%s:%d - %s", filepath.Base(file), line, content) } } l.handleLogEntry(entry) // 更新扫描状态 if level == LevelSuccess { l.scanStatus.UpdateSuccess() } else if level == LevelError { l.scanStatus.UpdateError() } } // shouldLog 检查是否应该记录该级别的日志 func (l *Logger) shouldLog(level LogLevel) bool { switch l.config.Level { case LevelAll: return true case LevelBaseInfoSuccess: return level == LevelBase || level == LevelInfo || level == LevelSuccess case LevelInfoSuccess: return level == LevelInfo || level == LevelSuccess case LevelError: return level == LevelError case LevelBase: return level == LevelBase case LevelInfo: return level == LevelInfo case LevelSuccess: return level == LevelSuccess case LevelDebug: return level == LevelDebug default: // 向后兼容:如果是字符串 "debug",显示所有 if l.config.Level == "debug" { return true } // 默认显示base、info和success return level == LevelBase || level == LevelInfo || level == LevelSuccess } } // handleLogEntry 处理日志条目 func (l *Logger) handleLogEntry(entry *LogEntry) { l.outputMutex.Lock() defer l.outputMutex.Unlock() // 清除进度条 l.clearProgress() // 使用所有处理器处理日志 l.mu.RLock() for _, handler := range l.handlers { if handler.IsEnabled() { handler.Handle(entry) } } l.mu.RUnlock() // 恢复进度条 l.restoreProgress() } // clearProgress 清除进度条 func (l *Logger) clearProgress() { if l.progressBar != nil { l.progressBar.Clear() // 忽略错误 time.Sleep(10 * time.Millisecond) } } // restoreProgress 恢复进度条 func (l *Logger) restoreProgress() { if l.progressBar != nil { l.progressBar.RenderBlank() // 忽略错误 } } // 便利方法 func (l *Logger) Debug(content string, metadata ...map[string]interface{}) { l.Log(LevelDebug, content, metadata...) } func (l *Logger) Base(content string, metadata ...map[string]interface{}) { l.Log(LevelBase, content, metadata...) } func (l *Logger) Info(content string, metadata ...map[string]interface{}) { l.Log(LevelInfo, content, metadata...) } func (l *Logger) Success(content string, metadata ...map[string]interface{}) { l.Log(LevelSuccess, content, metadata...) } func (l *Logger) Error(content string, metadata ...map[string]interface{}) { l.Log(LevelError, content, metadata...) } // GetScanStatus 获取扫描状态管理器 func (l *Logger) GetScanStatus() *ScanStatus { return l.scanStatus } // Initialize 初始化日志系统(兼容原接口) func (l *Logger) Initialize() { // 禁用标准日志输出 log.SetOutput(io.Discard) } // ConsoleHandler 控制台日志处理器 type ConsoleHandler struct { config *LoggerConfig formatter LogFormatter enabled bool mu sync.RWMutex } // NewConsoleHandler 创建控制台处理器 func NewConsoleHandler(config *LoggerConfig) *ConsoleHandler { formatter := NewStandardFormatter() formatter.SetStartTime(config.StartTime) return &ConsoleHandler{ config: config, formatter: formatter, enabled: true, } } // Handle 处理日志条目 func (h *ConsoleHandler) Handle(entry *LogEntry) { h.mu.RLock() defer h.mu.RUnlock() if !h.enabled { return } // 使用自己的格式化器格式化消息 logMsg := h.formatter.Format(entry) // 根据颜色设置输出 if h.config.EnableColor { if colorAttr, ok := h.config.LevelColors[entry.Level]; ok { color.New(colorAttr).Println(logMsg) } else { fmt.Println(logMsg) } } else { fmt.Println(logMsg) } // 根据慢速输出设置决定是否添加延迟 if h.config.SlowOutput { time.Sleep(50 * time.Millisecond) } } // SetEnabled 设置处理器启用状态 func (h *ConsoleHandler) SetEnabled(enabled bool) { h.mu.Lock() defer h.mu.Unlock() h.enabled = enabled } // IsEnabled 检查处理器是否启用 func (h *ConsoleHandler) IsEnabled() bool { h.mu.RLock() defer h.mu.RUnlock() return h.enabled } // 错误检查函数(保持原有逻辑) func CheckErrs(err error) error { if err == nil { return nil } // 已知需要重试的错误列表 errs := []string{ "closed by the remote host", "too many connections", "EOF", "A connection attempt failed", "established connection failed", "connection attempt failed", "Unable to read", "is not allowed to connect to this", "no pg_hba.conf entry", "No connection could be made", "invalid packet size", "bad connection", } // 检查错误是否匹配 errLower := strings.ToLower(err.Error()) for _, key := range errs { if strings.Contains(errLower, strings.ToLower(key)) { time.Sleep(1 * time.Second) return err } } return nil }