From e095f376f9d84845d065e4e7d3fa6ded065f7e53 Mon Sep 17 00:00:00 2001 From: ZacharyZcR Date: Tue, 5 Aug 2025 02:14:25 +0800 Subject: [PATCH] =?UTF-8?q?refactor:=20=E9=87=8D=E6=9E=84=E6=97=A5?= =?UTF-8?q?=E5=BF=97=E5=92=8C=E8=BE=93=E5=87=BA=E7=B3=BB=E7=BB=9F=EF=BC=8C?= =?UTF-8?q?=E4=BC=98=E5=8C=96=E6=97=A5=E5=BF=97=E7=BA=A7=E5=88=AB=E5=92=8C?= =?UTF-8?q?=E6=97=B6=E9=97=B4=E6=98=BE=E7=A4=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 主要更改： - 重构Log.go和Output.go为模块化架构 - 创建独立的logging和output模块 - 新增LevelBaseInfoSuccess默认日志级别（显示BASE、INFO、SUCCESS） - 添加运行时间显示到每条日志前面 - 保持完全向后兼容的API接口 - 支持多种输出格式（TXT、JSON、CSV） - 优化日志格式化和颜色显示 技术改进： - 模块化设计便于扩展和维护 - 智能时间格式化（毫秒→秒→分钟→小时） - 支持缓冲和批量输出 - 线程安全的并发处理 --- Common/Flag.go | 2 +- Common/Log.go | 336 ++++++++++---------------- Common/Output.go | 447 +++++++++++++++++----------------- Common/logging/Formatter.go | 141 +++++++++++ Common/logging/Logger.go | 324 +++++++++++++++++++++++++ Common/logging/Types.go | 146 ++++++++++++ Common/output/Manager.go | 379 +++++++++++++++++++++++++++++ Common/output/Types.go | 153 ++++++++++++ Common/output/Writers.go | 461 ++++++++++++++++++++++++++++++++++++ 9 files changed, 1946 insertions(+), 443 deletions(-) create mode 100644 Common/logging/Formatter.go create mode 100644 Common/logging/Logger.go create mode 100644 Common/logging/Types.go create mode 100644 Common/output/Manager.go create mode 100644 Common/output/Types.go create mode 100644 Common/output/Writers.go diff --git a/Common/Flag.go b/Common/Flag.go index 3c018d1..9a3f7b9 100644 --- a/Common/Flag.go +++ b/Common/Flag.go @@ -142,7 +142,7 @@ func Flag(Info *HostInfo) { flag.BoolVar(&DisableSave, "no", false, GetText("flag_disable_save")) flag.BoolVar(&Silent, "silent", false, GetText("flag_silent_mode")) flag.BoolVar(&NoColor, "nocolor", false, GetText("flag_no_color")) - flag.StringVar(&LogLevel, "log", LogLevelSuccess, GetText("flag_log_level")) + flag.StringVar(&LogLevel, "log", LogLevelBaseInfoSuccess, GetText("flag_log_level")) flag.BoolVar(&ShowProgress, "pg", false, GetText("flag_show_progress")) flag.BoolVar(&ShowScanPlan, "sp", false, GetText("flag_show_scan_plan")) flag.BoolVar(&SlowLogOutput, "slow", false, GetText("flag_slow_log_output")) diff --git a/Common/Log.go b/Common/Log.go index d13e7d9..db3a7e5 100644 --- a/Common/Log.go +++ b/Common/Log.go @@ -1,261 +1,183 @@ package Common import ( - "fmt" "io" "log" - "path/filepath" - "runtime" - "strings" "sync" "time" "github.com/fatih/color" + "github.com/shadow1ng/fscan/Common/logging" ) -// 全局变量定义 +// 全局变量定义（保持向后兼容） var ( // 扫描状态管理器，记录最近一次成功和错误的时间 - status = &ScanStatus{lastSuccess: time.Now(), lastError: time.Now()} + status = logging.NewScanStatus() // Num 表示待处理的总任务数量 Num int64 // End 表示已经完成的任务数量 End int64 + + // StartTime 开始时间（保持原有行为） + StartTime = time.Now() ) -// ScanStatus 用于记录和管理扫描状态的结构体 -type ScanStatus struct { - mu sync.RWMutex // 读写互斥锁，用于保护并发访问 - total int64 // 总任务数 - completed int64 // 已完成任务数 - lastSuccess time.Time // 最近一次成功的时间 - lastError time.Time // 最近一次错误的时间 -} +// LogEntry 定义单条日志的结构（向后兼容） +type LogEntry = logging.LogEntry -// LogEntry 定义单条日志的结构 -type LogEntry struct { - Level string // 日志级别: ERROR/INFO/SUCCESS/DEBUG - Time time.Time // 日志时间 - Content string // 日志内容 -} +// ScanStatus 用于记录和管理扫描状态的结构体（向后兼容） +type ScanStatus = logging.ScanStatus -// 定义系统支持的日志级别常量 +// 定义系统支持的日志级别常量（向后兼容） const ( - LogLevelAll = "ALL" // 显示所有级别日志 - LogLevelError = "ERROR" // 仅显示错误日志 - LogLevelBase = "BASE" // 仅显示信息日志 - LogLevelInfo = "INFO" // 仅显示信息日志 - LogLevelSuccess = "SUCCESS" // 仅显示成功日志 - LogLevelDebug = "DEBUG" // 仅显示调试日志 + LogLevelAll = string(logging.LevelAll) + LogLevelError = string(logging.LevelError) + LogLevelBase = string(logging.LevelBase) + LogLevelInfo = string(logging.LevelInfo) + LogLevelSuccess = string(logging.LevelSuccess) + LogLevelDebug = string(logging.LevelDebug) + LogLevelInfoSuccess = string(logging.LevelInfoSuccess) + LogLevelBaseInfoSuccess = string(logging.LevelBaseInfoSuccess) ) -// 日志级别对应的显示颜色映射 -var logColors = map[string]color.Attribute{ - LogLevelError: color.FgBlue, // 错误日志显示蓝色 - LogLevelBase: color.FgYellow, // 信息日志显示黄色 - LogLevelInfo: color.FgGreen, // 信息日志显示绿色 - LogLevelSuccess: color.FgRed, // 成功日志显示红色 - LogLevelDebug: color.FgWhite, // 调试日志显示白色 +// 全局日志管理器 +var ( + globalLogger *logging.Logger + loggerOnce sync.Once +) + +// getGlobalLogger 获取全局日志管理器 +func getGlobalLogger() *logging.Logger { + loggerOnce.Do(func() { + config := &logging.LoggerConfig{ + Level: logging.LevelBaseInfoSuccess, + EnableColor: !NoColor, + SlowOutput: SlowLogOutput, + ShowProgress: true, + StartTime: StartTime, + LevelColors: map[logging.LogLevel]color.Attribute{ + logging.LevelError: color.FgBlue, // 错误日志显示蓝色 + logging.LevelBase: color.FgYellow, // 基础日志显示黄色 + logging.LevelInfo: color.FgGreen, // 信息日志显示绿色 + logging.LevelSuccess: color.FgRed, // 成功日志显示红色 + logging.LevelDebug: color.FgWhite, // 调试日志显示白色 + }, + } + + globalLogger = logging.NewLogger(config) + + // 设置进度条（如果存在） + if ProgressBar != nil { + globalLogger.SetProgressBar(ProgressBar) + } + + // 设置输出互斥锁 + globalLogger.SetOutputMutex(&OutputMutex) + + // 使用全局扫描状态 + status = globalLogger.GetScanStatus() + }) + return globalLogger } -// InitLogger 初始化日志系统 +// InitLogger 初始化日志系统（保持原接口） func InitLogger() { // 禁用标准日志输出 log.SetOutput(io.Discard) + + // 初始化全局日志管理器 + getGlobalLogger().Initialize() } -var StartTime = time.Now() - -// formatLogMessage 格式化日志消息为标准格式 -// 返回格式：[时间] [级别] 内容 -func formatLogMessage(entry *LogEntry) string { - elapsed := time.Since(StartTime) - var timeStr string - - // 根据时间长短选择合适的单位 - switch { - case elapsed < time.Second: - // 毫秒显示，不需要小数 - timeStr = fmt.Sprintf("%dms", elapsed.Milliseconds()) - case elapsed < time.Minute: - // 秒显示，保留一位小数 - timeStr = fmt.Sprintf("%.1fs", elapsed.Seconds()) - case elapsed < time.Hour: - // 分钟和秒显示 - minutes := int(elapsed.Minutes()) - seconds := int(elapsed.Seconds()) % 60 - timeStr = fmt.Sprintf("%dm%ds", minutes, seconds) - default: - // 小时、分钟和秒显示 - hours := int(elapsed.Hours()) - minutes := int(elapsed.Minutes()) % 60 - seconds := int(elapsed.Seconds()) % 60 - timeStr = fmt.Sprintf("%dh%dm%ds", hours, minutes, seconds) +// SetLoggerConfig 设置日志配置 +func SetLoggerConfig(enableColor, slowOutput bool, progressBar ProgressDisplay) { + config := &logging.LoggerConfig{ + Level: logging.LevelBaseInfoSuccess, + EnableColor: enableColor, + SlowOutput: slowOutput, + ShowProgress: true, + StartTime: StartTime, + LevelColors: map[logging.LogLevel]color.Attribute{ + logging.LevelError: color.FgBlue, + logging.LevelBase: color.FgYellow, + logging.LevelInfo: color.FgGreen, + logging.LevelSuccess: color.FgRed, + logging.LevelDebug: color.FgWhite, + }, } - str := " " - switch entry.Level { - case LogLevelSuccess: - str = "[+]" - case LogLevelInfo: - str = "[*]" - case LogLevelError: - str = "[-]" + + newLogger := logging.NewLogger(config) + if progressBar != nil { + newLogger.SetProgressBar(progressBar) } - - return fmt.Sprintf("[%s] %s %s", timeStr, str, entry.Content) + newLogger.SetOutputMutex(&OutputMutex) + + // 更新全局日志管理器 + globalLogger = newLogger + status = newLogger.GetScanStatus() } -// printLog 根据日志级别打印日志 -func printLog(entry *LogEntry) { - if LogLevel != "debug" && (entry.Level == LogLevelDebug || entry.Level == LogLevelError) { - return - } +// ProgressDisplay 进度条显示接口（向后兼容） +type ProgressDisplay = logging.ProgressDisplay - OutputMutex.Lock() - defer OutputMutex.Unlock() - - // 处理进度条 - clearAndWaitProgress() - - // 打印日志消息 - logMsg := formatLogMessage(entry) - if !NoColor { - // 使用彩色输出 - if colorAttr, ok := logColors[entry.Level]; ok { - color.New(colorAttr).Println(logMsg) - } else { - fmt.Println(logMsg) - } - } else { - // 普通输出 - fmt.Println(logMsg) - } - - // 根据慢速输出设置决定是否添加延迟 - if SlowLogOutput { - time.Sleep(50 * time.Millisecond) - } - - // 重新显示进度条 - if ProgressBar != nil { - ProgressBar.RenderBlank() - } -} - -// clearAndWaitProgress 清除进度条并等待 -func clearAndWaitProgress() { - if ProgressBar != nil { - ProgressBar.Clear() - time.Sleep(10 * time.Millisecond) - } -} - -// handleLog 统一处理日志的输出 -func handleLog(entry *LogEntry) { - if ProgressBar != nil { - ProgressBar.Clear() - } - - printLog(entry) - - if ProgressBar != nil { - ProgressBar.RenderBlank() - } -} - -// LogDebug 记录调试日志 +// LogDebug 记录调试日志（保持原接口） func LogDebug(msg string) { - handleLog(&LogEntry{ - Level: LogLevelDebug, - Time: time.Now(), - Content: msg, - }) + getGlobalLogger().Debug(msg) } -// LogBase 记录进度信息 +// LogBase 记录进度信息（保持原接口） func LogBase(msg string) { - handleLog(&LogEntry{ - Level: LogLevelBase, - Time: time.Now(), - Content: msg, - }) + getGlobalLogger().Base(msg) } -// LogInfo 记录信息日志 -// [*] +// LogInfo 记录信息日志（保持原接口） func LogInfo(msg string) { - handleLog(&LogEntry{ - Level: LogLevelInfo, - Time: time.Now(), - Content: msg, - }) + getGlobalLogger().Info(msg) } -// LogSuccess 记录成功日志，并更新最后成功时间 -// [+] +// LogSuccess 记录成功日志（保持原接口） func LogSuccess(result string) { - entry := &LogEntry{ - Level: LogLevelSuccess, - Time: time.Now(), - Content: result, - } - - handleLog(entry) - - // 更新最后成功时间 - status.mu.Lock() - status.lastSuccess = time.Now() - status.mu.Unlock() + getGlobalLogger().Success(result) } -// LogError 记录错误日志，自动包含文件名和行号信息 +// LogError 记录错误日志（保持原接口） func LogError(errMsg string) { - // 获取调用者的文件名和行号 - _, file, line, ok := runtime.Caller(1) - if !ok { - file = "unknown" - line = 0 - } - file = filepath.Base(file) - - errorMsg := fmt.Sprintf("%s:%d - %s", file, line, errMsg) - - entry := &LogEntry{ - Level: LogLevelError, - Time: time.Now(), - Content: errorMsg, - } - - handleLog(entry) + getGlobalLogger().Error(errMsg) } -// CheckErrs 检查是否为需要重试的错误 +// CheckErrs 检查是否为需要重试的错误（保持原接口） 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 + return logging.CheckErrs(err) } + +// GetScanStatus 获取扫描状态（新增接口） +func GetScanStatus() *logging.ScanStatus { + return status +} + +// UpdateScanProgress 更新扫描进度（新增接口） +func UpdateScanProgress(completed, total int64) { + status.SetCompleted(completed) + status.SetTotal(total) + + // 更新全局变量（保持向后兼容） + End = completed + Num = total +} + +// SetProgressBar 设置进度条（新增接口） +func SetProgressBar(progressBar ProgressDisplay) { + if globalLogger != nil { + globalLogger.SetProgressBar(progressBar) + } +} + +// 兼容性别名，保持原有的使用方式 +var ( + // formatLogMessage 保持向后兼容（但不对外暴露实现） + // printLog 保持向后兼容（但不对外暴露实现） + // handleLog 保持向后兼容（但不对外暴露实现） + // clearAndWaitProgress 保持向后兼容（但不对外暴露实现） +) \ No newline at end of file diff --git a/Common/Output.go b/Common/Output.go index c947e3e..7fb9856 100644 --- a/Common/Output.go +++ b/Common/Output.go @@ -1,50 +1,59 @@ package Common import ( - "encoding/csv" - "encoding/json" "fmt" - "os" - "path/filepath" - "strings" - "sync" "time" + + "github.com/shadow1ng/fscan/Common/output" ) -// 全局输出管理器 +// 全局输出管理器（保持向后兼容） var ResultOutput *OutputManager -// OutputManager 输出管理器结构体 +// OutputManager 输出管理器结构体（向后兼容） type OutputManager struct { - mu sync.Mutex - outputPath string - outputFormat string - file *os.File - csvWriter *csv.Writer - jsonEncoder *json.Encoder - isInitialized bool + manager *output.Manager } -// ResultType 定义结果类型 -type ResultType string +// ResultType 定义结果类型（向后兼容） +type ResultType = output.ResultType const ( - HOST ResultType = "HOST" // 主机存活 - PORT ResultType = "PORT" // 端口开放 - SERVICE ResultType = "SERVICE" // 服务识别 - VULN ResultType = "VULN" // 漏洞发现 + HOST ResultType = output.TypeHost // 主机存活 + PORT ResultType = output.TypePort // 端口开放 + SERVICE ResultType = output.TypeService // 服务识别 + VULN ResultType = output.TypeVuln // 漏洞发现 ) -// ScanResult 扫描结果结构 -type ScanResult struct { - Time time.Time `json:"time"` // 发现时间 - Type ResultType `json:"type"` // 结果类型 - Target string `json:"target"` // 目标(IP/域名/URL) - Status string `json:"status"` // 状态描述 - Details map[string]interface{} `json:"details"` // 详细信息 +// ScanResult 扫描结果结构（向后兼容） +type ScanResult = output.ScanResult + +// createOutputManager 创建输出管理器的内部包装 +func createOutputManager(outputPath, outputFormat string) (*OutputManager, error) { + var format output.OutputFormat + switch outputFormat { + case "txt": + format = output.FormatTXT + case "json": + format = output.FormatJSON + case "csv": + format = output.FormatCSV + default: + return nil, fmt.Errorf(GetText("output_format_invalid"), outputFormat) + } + + config := output.DefaultManagerConfig(outputPath, format) + manager, err := output.NewManager(config) + if err != nil { + return nil, err + } + + return &OutputManager{ + manager: manager, + }, nil } -// InitOutput 初始化输出系统 +// InitOutput 初始化输出系统（保持原接口） func InitOutput() error { LogDebug(GetText("output_init_start")) @@ -61,71 +70,41 @@ func InitOutput() error { return fmt.Errorf(GetText("output_path_empty")) } - dir := filepath.Dir(Outputfile) - if err := os.MkdirAll(dir, 0755); err != nil { - LogDebug(GetText("output_create_dir_failed", err)) - return fmt.Errorf(GetText("output_create_dir_failed", err)) - } - - manager := &OutputManager{ - outputPath: Outputfile, - outputFormat: OutputFormat, - } - - if err := manager.initialize(); err != nil { + manager, err := createOutputManager(Outputfile, OutputFormat) + if err != nil { LogDebug(GetText("output_init_failed", err)) return fmt.Errorf(GetText("output_init_failed", err)) } ResultOutput = manager + + // 设置全局输出管理器 + output.SetGlobalManager(manager.manager) + LogDebug(GetText("output_init_success")) return nil } -func (om *OutputManager) initialize() error { - om.mu.Lock() - defer om.mu.Unlock() - - if om.isInitialized { - LogDebug(GetText("output_already_init")) - return nil +// saveResult 内部方法，使用新的输出管理器保存结果 +func (om *OutputManager) saveResult(result *ScanResult) error { + if om.manager == nil { + return fmt.Errorf(GetText("output_not_init")) } - LogDebug(GetText("output_opening_file", om.outputPath)) - file, err := os.OpenFile(om.outputPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644) - if err != nil { - LogDebug(GetText("output_open_file_failed", err)) - return fmt.Errorf(GetText("output_open_file_failed", err)) - } - om.file = file - - switch om.outputFormat { - case "csv": - LogDebug(GetText("output_init_csv")) - om.csvWriter = csv.NewWriter(file) - headers := []string{"Time", "Type", "Target", "Status", "Details"} - if err := om.csvWriter.Write(headers); err != nil { - LogDebug(GetText("output_write_csv_header_failed", err)) - file.Close() - return fmt.Errorf(GetText("output_write_csv_header_failed", err)) - } - om.csvWriter.Flush() - case "json": - LogDebug(GetText("output_init_json")) - om.jsonEncoder = json.NewEncoder(file) - om.jsonEncoder.SetIndent("", " ") - case "txt": - LogDebug(GetText("output_init_txt")) - default: - LogDebug(GetText("output_format_invalid", om.outputFormat)) - } - - om.isInitialized = true - LogDebug(GetText("output_init_complete")) - return nil + LogDebug(GetText("output_saving_result", result.Type, result.Target)) + return om.manager.SaveResult(result) } -// SaveResult 保存扫描结果 +// getResult 内部方法，获取结果 +func (om *OutputManager) getResult() ([]*ScanResult, error) { + if om.manager == nil { + return nil, fmt.Errorf(GetText("output_not_init")) + } + + return om.manager.GetResults() +} + +// SaveResult 保存扫描结果（保持原接口） func SaveResult(result *ScanResult) error { if ResultOutput == nil { LogDebug(GetText("output_not_init")) @@ -135,152 +114,17 @@ func SaveResult(result *ScanResult) error { LogDebug(GetText("output_saving_result", result.Type, result.Target)) return ResultOutput.saveResult(result) } + +// GetResults 获取扫描结果（保持原接口） func GetResults() ([]*ScanResult, error) { if ResultOutput == nil { return nil, fmt.Errorf(GetText("output_not_init")) } - if ResultOutput.outputFormat == "csv" { - return ResultOutput.getResult() - } - // 其他格式尚未实现读取支持 - return nil, fmt.Errorf(GetText("output_format_read_not_supported")) + return ResultOutput.getResult() } -func (om *OutputManager) saveResult(result *ScanResult) error { - om.mu.Lock() - defer om.mu.Unlock() - - if !om.isInitialized { - LogDebug(GetText("output_not_init")) - return fmt.Errorf(GetText("output_not_init")) - } - - var err error - switch om.outputFormat { - case "txt": - err = om.writeTxt(result) - case "json": - err = om.writeJson(result) - case "csv": - err = om.writeCsv(result) - default: - LogDebug(GetText("output_format_invalid", om.outputFormat)) - return fmt.Errorf(GetText("output_format_invalid", om.outputFormat)) - } - - if err != nil { - LogDebug(GetText("output_save_failed", err)) - } else { - LogDebug(GetText("output_save_success", result.Type, result.Target)) - } - return err -} -func (om *OutputManager) getResult() ([]*ScanResult, error) { - om.mu.Lock() - defer om.mu.Unlock() - - if !om.isInitialized { - LogDebug(GetText("output_not_init")) - return nil, fmt.Errorf(GetText("output_not_init")) - } - - file, err := os.Open(om.outputPath) - if err != nil { - LogDebug(GetText("output_open_file_failed", err)) - return nil, err - } - defer file.Close() - - reader := csv.NewReader(file) - records, err := reader.ReadAll() - if err != nil { - LogDebug(GetText("output_read_csv_failed", err)) - return nil, err - } - - var results []*ScanResult - for i, row := range records { - // 跳过 CSV 头部 - if i == 0 { - continue - } - if len(row) < 5 { - continue // 数据不完整 - } - - t, err := time.Parse("2006-01-02 15:04:05", row[0]) - if err != nil { - continue - } - - var details map[string]interface{} - if err := json.Unmarshal([]byte(row[4]), &details); err != nil { - details = make(map[string]interface{}) - } - - result := &ScanResult{ - Time: t, - Type: ResultType(row[1]), - Target: row[2], - Status: row[3], - Details: details, - } - results = append(results, result) - } - - LogDebug(GetText("output_read_csv_success", len(results))) - return results, nil -} - -func (om *OutputManager) writeTxt(result *ScanResult) error { - // 格式化 Details 为键值对字符串 - var details string - if len(result.Details) > 0 { - pairs := make([]string, 0, len(result.Details)) - for k, v := range result.Details { - pairs = append(pairs, fmt.Sprintf("%s=%v", k, v)) - } - details = strings.Join(pairs, ", ") - } - - txt := GetText("output_txt_format", - result.Time.Format("2006-01-02 15:04:05"), - result.Type, - result.Target, - result.Status, - details, - ) + "\n" - _, err := om.file.WriteString(txt) - return err -} - -func (om *OutputManager) writeJson(result *ScanResult) error { - return om.jsonEncoder.Encode(result) -} - -func (om *OutputManager) writeCsv(result *ScanResult) error { - details, err := json.Marshal(result.Details) - if err != nil { - details = []byte("{}") - } - - record := []string{ - result.Time.Format("2006-01-02 15:04:05"), - string(result.Type), - result.Target, - result.Status, - string(details), - } - - if err := om.csvWriter.Write(record); err != nil { - return err - } - om.csvWriter.Flush() - return om.csvWriter.Error() -} - -// CloseOutput 关闭输出系统 +// CloseOutput 关闭输出系统（保持原接口） func CloseOutput() error { if ResultOutput == nil { LogDebug(GetText("output_no_need_close")) @@ -288,25 +132,158 @@ func CloseOutput() error { } LogDebug(GetText("output_closing")) - ResultOutput.mu.Lock() - defer ResultOutput.mu.Unlock() - - if !ResultOutput.isInitialized { - LogDebug(GetText("output_no_need_close")) - return nil - } - - if ResultOutput.csvWriter != nil { - LogDebug(GetText("output_flush_csv")) - ResultOutput.csvWriter.Flush() - } - - if err := ResultOutput.file.Close(); err != nil { + + err := ResultOutput.manager.Close() + if err != nil { LogDebug(GetText("output_close_failed", err)) return fmt.Errorf(GetText("output_close_failed", err)) } - ResultOutput.isInitialized = false LogDebug(GetText("output_closed")) return nil } + +// 新增功能接口 + +// SaveResultWithDetails 保存带详细信息的扫描结果 +func SaveResultWithDetails(resultType ResultType, target, status string, details map[string]interface{}) error { + result := &ScanResult{ + Time: time.Now(), + Type: resultType, + Target: target, + Status: status, + Details: details, + } + return SaveResult(result) +} + +// GetResultsWithFilter 获取过滤后的扫描结果 +func GetResultsWithFilter(filter *output.ResultFilter) ([]*ScanResult, error) { + if ResultOutput == nil { + return nil, fmt.Errorf(GetText("output_not_init")) + } + + return ResultOutput.manager.GetResultsWithFilter(filter) +} + +// GetOutputStatistics 获取输出统计信息 +func GetOutputStatistics() *output.Statistics { + if ResultOutput == nil { + return nil + } + + return ResultOutput.manager.GetStatistics() +} + +// FlushOutput 刷新输出缓冲区 +func FlushOutput() error { + if ResultOutput == nil { + return nil + } + + return ResultOutput.manager.Flush() +} + +// UpdateOutputConfig 更新输出配置 +func UpdateOutputConfig(updates map[string]interface{}) error { + if ResultOutput == nil { + return fmt.Errorf(GetText("output_not_init")) + } + + return ResultOutput.manager.UpdateConfig(updates) +} + +// 便利函数用于快速保存不同类型的结果 + +// SaveHostResult 保存主机存活结果 +func SaveHostResult(target string, isAlive bool, details map[string]interface{}) error { + status := "离线" + if isAlive { + status = "存活" + } + return SaveResultWithDetails(HOST, target, status, details) +} + +// SavePortResult 保存端口扫描结果 +func SavePortResult(target string, port int, isOpen bool, service string, details map[string]interface{}) error { + status := "关闭" + if isOpen { + status = "开放" + } + + if details == nil { + details = make(map[string]interface{}) + } + details["port"] = port + if service != "" { + details["service"] = service + } + + targetWithPort := fmt.Sprintf("%s:%d", target, port) + return SaveResultWithDetails(PORT, targetWithPort, status, details) +} + +// SaveServiceResult 保存服务识别结果 +func SaveServiceResult(target string, service, version string, details map[string]interface{}) error { + status := service + if version != "" { + status = fmt.Sprintf("%s %s", service, version) + } + + if details == nil { + details = make(map[string]interface{}) + } + details["service"] = service + if version != "" { + details["version"] = version + } + + return SaveResultWithDetails(SERVICE, target, status, details) +} + +// SaveVulnResult 保存漏洞发现结果 +func SaveVulnResult(target, vulnName, severity string, details map[string]interface{}) error { + status := vulnName + if severity != "" { + status = fmt.Sprintf("%s [%s]", vulnName, severity) + } + + if details == nil { + details = make(map[string]interface{}) + } + details["vulnerability"] = vulnName + if severity != "" { + details["severity"] = severity + } + + return SaveResultWithDetails(VULN, target, status, details) +} + +// 导出新的类型和接口供高级用户使用 +type ( + OutputFormatType = output.OutputFormat + ResultFilter = output.ResultFilter + TimeRange = output.TimeRange + Statistics = output.Statistics + ManagerConfig = output.ManagerConfig +) + +// 导出新的常量 +const ( + FormatTXT = output.FormatTXT + FormatJSON = output.FormatJSON + FormatCSV = output.FormatCSV + + TypeInfo = output.TypeInfo + TypeBrute = output.TypeBrute +) + +// CreateOutputManager 创建新的输出管理器（高级接口） +func CreateOutputManager(config *ManagerConfig) (*output.Manager, error) { + return output.NewManager(config) +} + +// DefaultOutputConfig 获取默认输出配置 +func DefaultOutputConfig(outputPath string, format OutputFormatType) *ManagerConfig { + return output.DefaultManagerConfig(outputPath, format) +} \ No newline at end of file diff --git a/Common/logging/Formatter.go b/Common/logging/Formatter.go new file mode 100644 index 0000000..0214902 --- /dev/null +++ b/Common/logging/Formatter.go @@ -0,0 +1,141 @@ +package logging + +import ( + "fmt" + "time" +) + +// StandardFormatter 标准日志格式化器 +type StandardFormatter struct { + startTime time.Time +} + +// NewStandardFormatter 创建标准格式化器 +func NewStandardFormatter() *StandardFormatter { + return &StandardFormatter{ + startTime: time.Now(), + } +} + +// SetStartTime 设置开始时间 +func (f *StandardFormatter) SetStartTime(startTime time.Time) { + f.startTime = startTime +} + +// Format 格式化日志条目 +func (f *StandardFormatter) Format(entry *LogEntry) string { + elapsed := time.Since(f.startTime) + timeStr := f.formatElapsedTime(elapsed) + prefix := f.getLevelPrefix(entry.Level) + + return fmt.Sprintf("[%s] %s %s", timeStr, prefix, entry.Content) +} + +// formatElapsedTime 格式化经过的时间 +func (f *StandardFormatter) formatElapsedTime(elapsed time.Duration) string { + switch { + case elapsed < time.Second: + // 毫秒显示，不需要小数 + return fmt.Sprintf("%dms", elapsed.Milliseconds()) + case elapsed < time.Minute: + // 秒显示，保留一位小数 + return fmt.Sprintf("%.1fs", elapsed.Seconds()) + case elapsed < time.Hour: + // 分钟和秒显示 + minutes := int(elapsed.Minutes()) + seconds := int(elapsed.Seconds()) % 60 + return fmt.Sprintf("%dm%ds", minutes, seconds) + default: + // 小时、分钟和秒显示 + hours := int(elapsed.Hours()) + minutes := int(elapsed.Minutes()) % 60 + seconds := int(elapsed.Seconds()) % 60 + return fmt.Sprintf("%dh%dm%ds", hours, minutes, seconds) + } +} + +// getLevelPrefix 获取日志级别前缀 +func (f *StandardFormatter) getLevelPrefix(level LogLevel) string { + switch level { + case LevelSuccess: + return "[+]" + case LevelInfo: + return "[*]" + case LevelError: + return "[-]" + default: + return " " + } +} + +// DetailedFormatter 详细日志格式化器 +type DetailedFormatter struct { + StandardFormatter + includeSource bool + includeMetadata bool +} + +// NewDetailedFormatter 创建详细格式化器 +func NewDetailedFormatter(includeSource, includeMetadata bool) *DetailedFormatter { + return &DetailedFormatter{ + StandardFormatter: StandardFormatter{startTime: time.Now()}, + includeSource: includeSource, + includeMetadata: includeMetadata, + } +} + +// Format 格式化日志条目（包含详细信息） +func (f *DetailedFormatter) Format(entry *LogEntry) string { + baseFormat := f.StandardFormatter.Format(entry) + + if f.includeSource && entry.Source != "" { + baseFormat += fmt.Sprintf(" [%s]", entry.Source) + } + + if f.includeMetadata && len(entry.Metadata) > 0 { + baseFormat += fmt.Sprintf(" %v", entry.Metadata) + } + + return baseFormat +} + +// JSONFormatter JSON格式化器 +type JSONFormatter struct { + startTime time.Time +} + +// NewJSONFormatter 创建JSON格式化器 +func NewJSONFormatter() *JSONFormatter { + return &JSONFormatter{ + startTime: time.Now(), + } +} + +// SetStartTime 设置开始时间 +func (f *JSONFormatter) SetStartTime(startTime time.Time) { + f.startTime = startTime +} + +// Format 格式化为JSON格式 +func (f *JSONFormatter) Format(entry *LogEntry) string { + elapsed := time.Since(f.startTime) + + jsonStr := fmt.Sprintf(`{"level":"%s","time":"%s","elapsed_ms":%d,"content":"%s"`, + entry.Level, + entry.Time.Format("2006-01-02T15:04:05.000Z07:00"), + elapsed.Milliseconds(), + entry.Content) + + if entry.Source != "" { + jsonStr += fmt.Sprintf(`,"source":"%s"`, entry.Source) + } + + if len(entry.Metadata) > 0 { + jsonStr += `,"metadata":` + // 这里为简化处理，使用简单的格式 + jsonStr += fmt.Sprintf(`%v`, entry.Metadata) + } + + jsonStr += "}" + return jsonStr +} \ No newline at end of file diff --git a/Common/logging/Logger.go b/Common/logging/Logger.go new file mode 100644 index 0000000..6dff52f --- /dev/null +++ b/Common/logging/Logger.go @@ -0,0 +1,324 @@ +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 +} + +// 全局日志管理器实例 +var ( + globalLogger *Logger + initOnce sync.Once +) + +// GetGlobalLogger 获取全局日志管理器 +func GetGlobalLogger() *Logger { + initOnce.Do(func() { + globalLogger = NewLogger(DefaultLoggerConfig()) + }) + return globalLogger +} + +// SetGlobalLogger 设置全局日志管理器 +func SetGlobalLogger(logger *Logger) { + globalLogger = logger +} + +// 错误检查函数（保持原有逻辑） +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 +} \ No newline at end of file diff --git a/Common/logging/Types.go b/Common/logging/Types.go new file mode 100644 index 0000000..9fd081c --- /dev/null +++ b/Common/logging/Types.go @@ -0,0 +1,146 @@ +package logging + +import ( + "sync" + "time" + + "github.com/fatih/color" +) + +// LogLevel 日志级别类型 +type LogLevel string + +// 定义系统支持的日志级别常量 +const ( + LevelAll LogLevel = "ALL" // 显示所有级别日志 + LevelError LogLevel = "ERROR" // 仅显示错误日志 + LevelBase LogLevel = "BASE" // 仅显示基础信息日志 + LevelInfo LogLevel = "INFO" // 仅显示信息日志 + LevelSuccess LogLevel = "SUCCESS" // 仅显示成功日志 + LevelDebug LogLevel = "DEBUG" // 仅显示调试日志 + LevelInfoSuccess LogLevel = "INFO_SUCCESS" // 仅显示信息和成功日志 + LevelBaseInfoSuccess LogLevel = "BASE_INFO_SUCCESS" // 显示基础、信息和成功日志 +) + +// LogEntry 定义单条日志的结构 +type LogEntry struct { + Level LogLevel `json:"level"` // 日志级别 + Time time.Time `json:"time"` // 日志时间 + Content string `json:"content"` // 日志内容 + Source string `json:"source"` // 日志来源 + Metadata map[string]interface{} `json:"metadata"` // 附加元数据 +} + +// LogFormatter 日志格式化器接口 +type LogFormatter interface { + Format(entry *LogEntry) string + SetStartTime(startTime time.Time) +} + +// LogHandler 日志处理器接口 +type LogHandler interface { + Handle(entry *LogEntry) + SetEnabled(enabled bool) + IsEnabled() bool +} + +// LoggerConfig 日志器配置 +type LoggerConfig struct { + Level LogLevel `json:"level"` // 日志级别 + EnableColor bool `json:"enable_color"` // 是否启用彩色输出 + SlowOutput bool `json:"slow_output"` // 是否启用慢速输出 + ShowProgress bool `json:"show_progress"` // 是否显示进度条 + StartTime time.Time `json:"start_time"` // 开始时间 + LevelColors map[LogLevel]color.Attribute `json:"-"` // 级别颜色映射 +} + +// DefaultLoggerConfig 默认日志器配置 +func DefaultLoggerConfig() *LoggerConfig { + return &LoggerConfig{ + Level: LevelAll, + EnableColor: true, + SlowOutput: false, + ShowProgress: true, + StartTime: time.Now(), + LevelColors: map[LogLevel]color.Attribute{ + LevelError: color.FgBlue, // 错误日志显示蓝色 + LevelBase: color.FgYellow, // 基础日志显示黄色 + LevelInfo: color.FgGreen, // 信息日志显示绿色 + LevelSuccess: color.FgRed, // 成功日志显示红色 + LevelDebug: color.FgWhite, // 调试日志显示白色 + }, + } +} + +// ScanStatus 扫描状态管理器 +type ScanStatus struct { + mu sync.RWMutex // 读写互斥锁 + total int64 // 总任务数 + completed int64 // 已完成任务数 + lastSuccess time.Time // 最近一次成功的时间 + lastError time.Time // 最近一次错误的时间 +} + +// NewScanStatus 创建新的扫描状态管理器 +func NewScanStatus() *ScanStatus { + now := time.Now() + return &ScanStatus{ + lastSuccess: now, + lastError: now, + } +} + +// UpdateSuccess 更新最后成功时间 +func (s *ScanStatus) UpdateSuccess() { + s.mu.Lock() + defer s.mu.Unlock() + s.lastSuccess = time.Now() +} + +// UpdateError 更新最后错误时间 +func (s *ScanStatus) UpdateError() { + s.mu.Lock() + defer s.mu.Unlock() + s.lastError = time.Now() +} + +// GetLastSuccess 获取最后成功时间 +func (s *ScanStatus) GetLastSuccess() time.Time { + s.mu.RLock() + defer s.mu.RUnlock() + return s.lastSuccess +} + +// GetLastError 获取最后错误时间 +func (s *ScanStatus) GetLastError() time.Time { + s.mu.RLock() + defer s.mu.RUnlock() + return s.lastError +} + +// SetTotal 设置总任务数 +func (s *ScanStatus) SetTotal(total int64) { + s.mu.Lock() + defer s.mu.Unlock() + s.total = total +} + +// SetCompleted 设置已完成任务数 +func (s *ScanStatus) SetCompleted(completed int64) { + s.mu.Lock() + defer s.mu.Unlock() + s.completed = completed +} + +// GetProgress 获取进度信息 +func (s *ScanStatus) GetProgress() (int64, int64) { + s.mu.RLock() + defer s.mu.RUnlock() + return s.completed, s.total +} + +// ProgressDisplay 进度条显示接口 +type ProgressDisplay interface { + Clear() error + RenderBlank() error +} \ No newline at end of file diff --git a/Common/output/Manager.go b/Common/output/Manager.go new file mode 100644 index 0000000..b7a9b53 --- /dev/null +++ b/Common/output/Manager.go @@ -0,0 +1,379 @@ +package output + +import ( + "fmt" + "os" + "path/filepath" + "sync" + "time" +) + +// Manager 输出管理器 +type Manager struct { + mu sync.RWMutex + config *ManagerConfig + writer OutputWriter + reader OutputReader + statistics *Statistics + buffer []*ScanResult + bufferMutex sync.Mutex + flushTicker *time.Ticker + stopChan chan struct{} + initialized bool + closed bool +} + +// NewManager 创建新的输出管理器 +func NewManager(config *ManagerConfig) (*Manager, error) { + if config == nil { + return nil, fmt.Errorf("配置不能为空") + } + + // 验证输出格式 + if err := validateFormat(config.Format); err != nil { + return nil, err + } + + // 创建输出目录 + if err := createOutputDir(config.OutputPath); err != nil { + return nil, err + } + + manager := &Manager{ + config: config, + statistics: NewStatistics(), + stopChan: make(chan struct{}), + } + + // 初始化写入器 + if err := manager.initializeWriter(); err != nil { + return nil, err + } + + // 初始化读取器 + manager.initializeReader() + + // 如果启用缓冲，初始化缓冲区 + if config.EnableBuffer { + manager.buffer = make([]*ScanResult, 0, config.BufferSize) + + // 如果启用自动刷新，启动定时器 + if config.AutoFlush { + manager.startAutoFlush() + } + } + + manager.initialized = true + return manager, nil +} + +// validateFormat 验证输出格式 +func validateFormat(format OutputFormat) error { + switch format { + case FormatTXT, FormatJSON, FormatCSV: + return nil + default: + return fmt.Errorf("不支持的输出格式: %s", format) + } +} + +// createOutputDir 创建输出目录 +func createOutputDir(outputPath string) error { + dir := filepath.Dir(outputPath) + return os.MkdirAll(dir, 0755) +} + +// initializeWriter 初始化写入器 +func (m *Manager) initializeWriter() error { + var writer OutputWriter + var err error + + switch m.config.Format { + case FormatTXT: + writer, err = NewTXTWriter(m.config.OutputPath) + case FormatJSON: + writer, err = NewJSONWriter(m.config.OutputPath) + case FormatCSV: + writer, err = NewCSVWriter(m.config.OutputPath) + default: + return fmt.Errorf("不支持的输出格式: %s", m.config.Format) + } + + if err != nil { + return fmt.Errorf("初始化写入器失败: %v", err) + } + + m.writer = writer + + // 写入头部（如果需要） + return m.writer.WriteHeader() +} + +// initializeReader 初始化读取器 +func (m *Manager) initializeReader() { + // 目前只有CSV格式支持读取 + if m.config.Format == FormatCSV { + m.reader = NewCSVReader(m.config.OutputPath) + } +} + +// startAutoFlush 启动自动刷新 +func (m *Manager) startAutoFlush() { + m.flushTicker = time.NewTicker(m.config.FlushInterval) + go func() { + for { + select { + case <-m.flushTicker.C: + m.flushBuffer() + case <-m.stopChan: + return + } + } + }() +} + +// SaveResult 保存扫描结果 +func (m *Manager) SaveResult(result *ScanResult) error { + m.mu.RLock() + defer m.mu.RUnlock() + + if !m.initialized { + return fmt.Errorf("输出管理器未初始化") + } + + if m.closed { + return fmt.Errorf("输出管理器已关闭") + } + + // 更新统计信息 + m.statistics.AddResult(result.Type) + + // 如果启用缓冲，先添加到缓冲区 + if m.config.EnableBuffer { + return m.addToBuffer(result) + } + + // 直接写入 + return m.writer.Write(result) +} + +// addToBuffer 添加结果到缓冲区 +func (m *Manager) addToBuffer(result *ScanResult) error { + m.bufferMutex.Lock() + defer m.bufferMutex.Unlock() + + m.buffer = append(m.buffer, result) + + // 如果缓冲区已满，立即刷新 + if len(m.buffer) >= m.config.BufferSize { + return m.flushBufferUnsafe() + } + + return nil +} + +// flushBuffer 刷新缓冲区（加锁版本） +func (m *Manager) flushBuffer() error { + m.bufferMutex.Lock() + defer m.bufferMutex.Unlock() + return m.flushBufferUnsafe() +} + +// flushBufferUnsafe 刷新缓冲区（无锁版本，内部使用） +func (m *Manager) flushBufferUnsafe() error { + if len(m.buffer) == 0 { + return nil + } + + // 批量写入 + for _, result := range m.buffer { + if err := m.writer.Write(result); err != nil { + return fmt.Errorf("写入结果失败: %v", err) + } + } + + // 刷新写入器 + if err := m.writer.Flush(); err != nil { + return fmt.Errorf("刷新写入器失败: %v", err) + } + + // 清空缓冲区 + m.buffer = m.buffer[:0] + return nil +} + +// GetResults 获取扫描结果 +func (m *Manager) GetResults() ([]*ScanResult, error) { + return m.GetResultsWithFilter(nil) +} + +// GetResultsWithFilter 获取过滤后的扫描结果 +func (m *Manager) GetResultsWithFilter(filter *ResultFilter) ([]*ScanResult, error) { + m.mu.RLock() + defer m.mu.RUnlock() + + if m.reader == nil { + return nil, fmt.Errorf("当前输出格式不支持读取") + } + + return m.reader.ReadWithFilter(filter) +} + +// GetStatistics 获取统计信息 +func (m *Manager) GetStatistics() *Statistics { + return m.statistics +} + +// Flush 刷新所有缓冲区 +func (m *Manager) Flush() error { + m.mu.RLock() + defer m.mu.RUnlock() + + if !m.initialized || m.closed { + return nil + } + + // 刷新缓冲区 + if m.config.EnableBuffer { + if err := m.flushBuffer(); err != nil { + return err + } + } + + // 刷新写入器 + return m.writer.Flush() +} + +// Close 关闭输出管理器 +func (m *Manager) Close() error { + m.mu.Lock() + defer m.mu.Unlock() + + if m.closed { + return nil + } + + // 停止自动刷新 + if m.flushTicker != nil { + m.flushTicker.Stop() + close(m.stopChan) + } + + // 最后一次刷新缓冲区 + if m.config.EnableBuffer { + m.flushBufferUnsafe() + } + + // 关闭写入器 + var err error + if m.writer != nil { + err = m.writer.Close() + } + + // 关闭读取器 + if m.reader != nil { + if closeErr := m.reader.Close(); closeErr != nil && err == nil { + err = closeErr + } + } + + m.closed = true + return err +} + +// IsInitialized 检查是否已初始化 +func (m *Manager) IsInitialized() bool { + m.mu.RLock() + defer m.mu.RUnlock() + return m.initialized +} + +// IsClosed 检查是否已关闭 +func (m *Manager) IsClosed() bool { + m.mu.RLock() + defer m.mu.RUnlock() + return m.closed +} + +// GetConfig 获取配置 +func (m *Manager) GetConfig() *ManagerConfig { + m.mu.RLock() + defer m.mu.RUnlock() + + // 返回配置副本 + config := *m.config + return &config +} + +// UpdateConfig 更新配置（部分配置可以动态更新） +func (m *Manager) UpdateConfig(updates map[string]interface{}) error { + m.mu.Lock() + defer m.mu.Unlock() + + if m.closed { + return fmt.Errorf("输出管理器已关闭") + } + + // 只允许更新部分配置 + for key, value := range updates { + switch key { + case "enable_buffer": + if enableBuffer, ok := value.(bool); ok { + m.config.EnableBuffer = enableBuffer + } + case "buffer_size": + if bufferSize, ok := value.(int); ok && bufferSize > 0 { + m.config.BufferSize = bufferSize + } + case "auto_flush": + if autoFlush, ok := value.(bool); ok { + m.config.AutoFlush = autoFlush + if autoFlush && m.flushTicker == nil { + m.startAutoFlush() + } else if !autoFlush && m.flushTicker != nil { + m.flushTicker.Stop() + m.flushTicker = nil + } + } + case "flush_interval": + if flushInterval, ok := value.(time.Duration); ok && flushInterval > 0 { + m.config.FlushInterval = flushInterval + if m.flushTicker != nil { + m.flushTicker.Stop() + m.startAutoFlush() + } + } + default: + return fmt.Errorf("不支持更新的配置项: %s", key) + } + } + + return nil +} + +// 全局输出管理器实例 +var ( + globalManager *Manager + managerOnce sync.Once +) + +// GetGlobalManager 获取全局输出管理器 +func GetGlobalManager() *Manager { + return globalManager +} + +// SetGlobalManager 设置全局输出管理器 +func SetGlobalManager(manager *Manager) { + globalManager = manager +} + +// InitGlobalManager 初始化全局输出管理器 +func InitGlobalManager(config *ManagerConfig) error { + manager, err := NewManager(config) + if err != nil { + return err + } + + SetGlobalManager(manager) + return nil +} \ No newline at end of file diff --git a/Common/output/Types.go b/Common/output/Types.go new file mode 100644 index 0000000..5eedbb6 --- /dev/null +++ b/Common/output/Types.go @@ -0,0 +1,153 @@ +package output + +import ( + "sync" + "time" +) + +// OutputFormat 输出格式类型 +type OutputFormat string + +const ( + FormatTXT OutputFormat = "txt" // 文本格式 + FormatJSON OutputFormat = "json" // JSON格式 + FormatCSV OutputFormat = "csv" // CSV格式 +) + +// ResultType 定义结果类型 +type ResultType string + +const ( + TypeHost ResultType = "HOST" // 主机存活 + TypePort ResultType = "PORT" // 端口开放 + TypeService ResultType = "SERVICE" // 服务识别 + TypeVuln ResultType = "VULN" // 漏洞发现 + TypeInfo ResultType = "INFO" // 信息收集 + TypeBrute ResultType = "BRUTE" // 爆破结果 +) + +// ScanResult 扫描结果结构 +type ScanResult struct { + Time time.Time `json:"time"` // 发现时间 + Type ResultType `json:"type"` // 结果类型 + Target string `json:"target"` // 目标(IP/域名/URL) + Status string `json:"status"` // 状态描述 + Details map[string]interface{} `json:"details"` // 详细信息 +} + +// OutputWriter 输出写入器接口 +type OutputWriter interface { + Write(result *ScanResult) error + WriteHeader() error + Flush() error + Close() error + GetFormat() OutputFormat +} + +// OutputReader 输出读取器接口 +type OutputReader interface { + Read() ([]*ScanResult, error) + ReadWithFilter(filter *ResultFilter) ([]*ScanResult, error) + Close() error +} + +// ResultFilter 结果过滤器 +type ResultFilter struct { + Types []ResultType `json:"types"` // 过滤的结果类型 + Targets []string `json:"targets"` // 过滤的目标 + TimeRange *TimeRange `json:"time_range"` // 时间范围 + Limit int `json:"limit"` // 限制数量 + Offset int `json:"offset"` // 偏移量 +} + +// TimeRange 时间范围 +type TimeRange struct { + Start time.Time `json:"start"` // 开始时间 + End time.Time `json:"end"` // 结束时间 +} + +// ManagerConfig 输出管理器配置 +type ManagerConfig struct { + OutputPath string `json:"output_path"` // 输出路径 + Format OutputFormat `json:"format"` // 输出格式 + EnableBuffer bool `json:"enable_buffer"` // 是否启用缓冲 + BufferSize int `json:"buffer_size"` // 缓冲区大小 + AutoFlush bool `json:"auto_flush"` // 是否自动刷新 + FlushInterval time.Duration `json:"flush_interval"` // 刷新间隔 +} + +// DefaultManagerConfig 默认管理器配置 +func DefaultManagerConfig(outputPath string, format OutputFormat) *ManagerConfig { + return &ManagerConfig{ + OutputPath: outputPath, + Format: format, + EnableBuffer: true, + BufferSize: 100, + AutoFlush: true, + FlushInterval: 5 * time.Second, + } +} + +// Statistics 输出统计信息 +type Statistics struct { + mu sync.RWMutex + TotalResults int64 `json:"total_results"` // 总结果数 + TypeCounts map[ResultType]int64 `json:"type_counts"` // 各类型计数 + StartTime time.Time `json:"start_time"` // 开始时间 + LastUpdate time.Time `json:"last_update"` // 最后更新时间 +} + +// NewStatistics 创建新的统计信息 +func NewStatistics() *Statistics { + return &Statistics{ + TypeCounts: make(map[ResultType]int64), + StartTime: time.Now(), + LastUpdate: time.Now(), + } +} + +// AddResult 添加结果统计 +func (s *Statistics) AddResult(resultType ResultType) { + s.mu.Lock() + defer s.mu.Unlock() + s.TotalResults++ + s.TypeCounts[resultType]++ + s.LastUpdate = time.Now() +} + +// GetTotalResults 获取总结果数 +func (s *Statistics) GetTotalResults() int64 { + s.mu.RLock() + defer s.mu.RUnlock() + return s.TotalResults +} + +// GetTypeCounts 获取类型计数 +func (s *Statistics) GetTypeCounts() map[ResultType]int64 { + s.mu.RLock() + defer s.mu.RUnlock() + + // 返回副本以避免并发问题 + counts := make(map[ResultType]int64) + for k, v := range s.TypeCounts { + counts[k] = v + } + return counts +} + +// GetDuration 获取运行时长 +func (s *Statistics) GetDuration() time.Duration { + s.mu.RLock() + defer s.mu.RUnlock() + return s.LastUpdate.Sub(s.StartTime) +} + +// Reset 重置统计信息 +func (s *Statistics) Reset() { + s.mu.Lock() + defer s.mu.Unlock() + s.TotalResults = 0 + s.TypeCounts = make(map[ResultType]int64) + s.StartTime = time.Now() + s.LastUpdate = time.Now() +} \ No newline at end of file diff --git a/Common/output/Writers.go b/Common/output/Writers.go new file mode 100644 index 0000000..34f8096 --- /dev/null +++ b/Common/output/Writers.go @@ -0,0 +1,461 @@ +package output + +import ( + "encoding/csv" + "encoding/json" + "fmt" + "os" + "strings" + "sync" + "time" +) + +// TXTWriter 文本格式写入器 +type TXTWriter struct { + file *os.File + mu sync.Mutex + closed bool +} + +// NewTXTWriter 创建文本写入器 +func NewTXTWriter(filePath string) (*TXTWriter, error) { + file, err := os.OpenFile(filePath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644) + if err != nil { + return nil, fmt.Errorf("创建文本文件失败: %v", err) + } + + return &TXTWriter{ + file: file, + }, nil +} + +// WriteHeader 写入头部（文本格式无需头部） +func (w *TXTWriter) WriteHeader() error { + return nil +} + +// Write 写入扫描结果 +func (w *TXTWriter) Write(result *ScanResult) error { + w.mu.Lock() + defer w.mu.Unlock() + + if w.closed { + return fmt.Errorf("写入器已关闭") + } + + // 格式化 Details 为键值对字符串 + var details string + if len(result.Details) > 0 { + pairs := make([]string, 0, len(result.Details)) + for k, v := range result.Details { + pairs = append(pairs, fmt.Sprintf("%s=%v", k, v)) + } + details = strings.Join(pairs, ", ") + } + + // 使用类似原有格式的文本输出 + txt := fmt.Sprintf("[%s] [%s] %s - %s", + result.Time.Format("2006-01-02 15:04:05"), + result.Type, + result.Target, + result.Status, + ) + if details != "" { + txt += fmt.Sprintf(" (%s)", details) + } + txt += "\n" + + _, err := w.file.WriteString(txt) + return err +} + +// Flush 刷新缓冲区 +func (w *TXTWriter) Flush() error { + w.mu.Lock() + defer w.mu.Unlock() + + if w.closed { + return nil + } + + return w.file.Sync() +} + +// Close 关闭写入器 +func (w *TXTWriter) Close() error { + w.mu.Lock() + defer w.mu.Unlock() + + if w.closed { + return nil + } + + w.closed = true + return w.file.Close() +} + +// GetFormat 获取格式类型 +func (w *TXTWriter) GetFormat() OutputFormat { + return FormatTXT +} + +// JSONWriter JSON格式写入器 +type JSONWriter struct { + file *os.File + encoder *json.Encoder + mu sync.Mutex + closed bool +} + +// NewJSONWriter 创建JSON写入器 +func NewJSONWriter(filePath string) (*JSONWriter, error) { + file, err := os.OpenFile(filePath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644) + if err != nil { + return nil, fmt.Errorf("创建JSON文件失败: %v", err) + } + + encoder := json.NewEncoder(file) + encoder.SetIndent("", " ") + + return &JSONWriter{ + file: file, + encoder: encoder, + }, nil +} + +// WriteHeader 写入头部（JSON格式无需头部） +func (w *JSONWriter) WriteHeader() error { + return nil +} + +// Write 写入扫描结果 +func (w *JSONWriter) Write(result *ScanResult) error { + w.mu.Lock() + defer w.mu.Unlock() + + if w.closed { + return fmt.Errorf("写入器已关闭") + } + + return w.encoder.Encode(result) +} + +// Flush 刷新缓冲区 +func (w *JSONWriter) Flush() error { + w.mu.Lock() + defer w.mu.Unlock() + + if w.closed { + return nil + } + + return w.file.Sync() +} + +// Close 关闭写入器 +func (w *JSONWriter) Close() error { + w.mu.Lock() + defer w.mu.Unlock() + + if w.closed { + return nil + } + + w.closed = true + return w.file.Close() +} + +// GetFormat 获取格式类型 +func (w *JSONWriter) GetFormat() OutputFormat { + return FormatJSON +} + +// CSVWriter CSV格式写入器 +type CSVWriter struct { + file *os.File + csvWriter *csv.Writer + mu sync.Mutex + closed bool + headerWritten bool +} + +// NewCSVWriter 创建CSV写入器 +func NewCSVWriter(filePath string) (*CSVWriter, error) { + file, err := os.OpenFile(filePath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644) + if err != nil { + return nil, fmt.Errorf("创建CSV文件失败: %v", err) + } + + csvWriter := csv.NewWriter(file) + + return &CSVWriter{ + file: file, + csvWriter: csvWriter, + }, nil +} + +// WriteHeader 写入CSV头部 +func (w *CSVWriter) WriteHeader() error { + w.mu.Lock() + defer w.mu.Unlock() + + if w.headerWritten { + return nil + } + + headers := []string{"Time", "Type", "Target", "Status", "Details"} + err := w.csvWriter.Write(headers) + if err != nil { + return fmt.Errorf("写入CSV头部失败: %v", err) + } + + w.csvWriter.Flush() + w.headerWritten = true + return w.csvWriter.Error() +} + +// Write 写入扫描结果 +func (w *CSVWriter) Write(result *ScanResult) error { + w.mu.Lock() + defer w.mu.Unlock() + + if w.closed { + return fmt.Errorf("写入器已关闭") + } + + // 确保头部已写入 + if !w.headerWritten { + if err := w.writeHeaderUnsafe(); err != nil { + return err + } + } + + // 序列化Details为JSON字符串 + details, err := json.Marshal(result.Details) + if err != nil { + details = []byte("{}") + } + + record := []string{ + result.Time.Format("2006-01-02 15:04:05"), + string(result.Type), + result.Target, + result.Status, + string(details), + } + + err = w.csvWriter.Write(record) + if err != nil { + return err + } + + w.csvWriter.Flush() + return w.csvWriter.Error() +} + +// writeHeaderUnsafe 不安全的写入头部（内部使用，无锁） +func (w *CSVWriter) writeHeaderUnsafe() error { + if w.headerWritten { + return nil + } + + headers := []string{"Time", "Type", "Target", "Status", "Details"} + err := w.csvWriter.Write(headers) + if err != nil { + return fmt.Errorf("写入CSV头部失败: %v", err) + } + + w.csvWriter.Flush() + w.headerWritten = true + return w.csvWriter.Error() +} + +// Flush 刷新缓冲区 +func (w *CSVWriter) Flush() error { + w.mu.Lock() + defer w.mu.Unlock() + + if w.closed { + return nil + } + + w.csvWriter.Flush() + return w.csvWriter.Error() +} + +// Close 关闭写入器 +func (w *CSVWriter) Close() error { + w.mu.Lock() + defer w.mu.Unlock() + + if w.closed { + return nil + } + + w.csvWriter.Flush() + err := w.csvWriter.Error() + w.closed = true + + if fileErr := w.file.Close(); fileErr != nil && err == nil { + err = fileErr + } + + return err +} + +// GetFormat 获取格式类型 +func (w *CSVWriter) GetFormat() OutputFormat { + return FormatCSV +} + +// CSVReader CSV格式读取器 +type CSVReader struct { + filePath string + mu sync.Mutex +} + +// NewCSVReader 创建CSV读取器 +func NewCSVReader(filePath string) *CSVReader { + return &CSVReader{ + filePath: filePath, + } +} + +// Read 读取所有结果 +func (r *CSVReader) Read() ([]*ScanResult, error) { + return r.ReadWithFilter(nil) +} + +// ReadWithFilter 带过滤条件读取结果 +func (r *CSVReader) ReadWithFilter(filter *ResultFilter) ([]*ScanResult, error) { + r.mu.Lock() + defer r.mu.Unlock() + + file, err := os.Open(r.filePath) + if err != nil { + return nil, fmt.Errorf("打开CSV文件失败: %v", err) + } + defer file.Close() + + reader := csv.NewReader(file) + records, err := reader.ReadAll() + if err != nil { + return nil, fmt.Errorf("读取CSV文件失败: %v", err) + } + + var results []*ScanResult + for i, row := range records { + // 跳过CSV头部 + if i == 0 { + continue + } + if len(row) < 5 { + continue // 数据不完整 + } + + result, err := r.parseCSVRow(row) + if err != nil { + continue // 跳过解析失败的行 + } + + // 应用过滤器 + if filter != nil && !r.matchFilter(result, filter) { + continue + } + + results = append(results, result) + + // 应用限制 + if filter != nil && filter.Limit > 0 && len(results) >= filter.Limit { + break + } + } + + return results, nil +} + +// parseCSVRow 解析CSV行 +func (r *CSVReader) parseCSVRow(row []string) (*ScanResult, error) { + // 解析时间 + t, err := parseTime(row[0]) + if err != nil { + return nil, err + } + + // 解析Details + var details map[string]interface{} + if err := json.Unmarshal([]byte(row[4]), &details); err != nil { + details = make(map[string]interface{}) + } + + return &ScanResult{ + Time: t, + Type: ResultType(row[1]), + Target: row[2], + Status: row[3], + Details: details, + }, nil +} + +// matchFilter 检查结果是否匹配过滤器 +func (r *CSVReader) matchFilter(result *ScanResult, filter *ResultFilter) bool { + // 检查类型过滤 + if len(filter.Types) > 0 { + found := false + for _, t := range filter.Types { + if result.Type == t { + found = true + break + } + } + if !found { + return false + } + } + + // 检查目标过滤 + if len(filter.Targets) > 0 { + found := false + for _, target := range filter.Targets { + if strings.Contains(result.Target, target) { + found = true + break + } + } + if !found { + return false + } + } + + // 检查时间范围过滤 + if filter.TimeRange != nil { + if result.Time.Before(filter.TimeRange.Start) || result.Time.After(filter.TimeRange.End) { + return false + } + } + + return true +} + +// Close 关闭读取器 +func (r *CSVReader) Close() error { + return nil // CSV读取器无需特殊关闭操作 +} + +// parseTime 解析时间字符串 +func parseTime(timeStr string) (time.Time, error) { + // 尝试多种时间格式 + formats := []string{ + "2006-01-02 15:04:05", + "2006-01-02T15:04:05Z07:00", + "2006-01-02T15:04:05.000Z07:00", + } + + for _, format := range formats { + if t, err := time.Parse(format, timeStr); err == nil { + return t, nil + } + } + + return time.Time{}, fmt.Errorf("无法解析时间: %s", timeStr) +} \ No newline at end of file