From b463984e78e6ac2a3b2fe79d7eade868160a5ec7 Mon Sep 17 00:00:00 2001 From: ZacharyZcR Date: Tue, 12 Aug 2025 14:37:28 +0800 Subject: [PATCH] =?UTF-8?q?refactor:=20=E9=87=8D=E6=9E=84main.go=E5=85=A5?= =?UTF-8?q?=E5=8F=A3=E6=96=87=E4=BB=B6=E5=B9=B6=E5=BC=95=E5=85=A5=E4=BE=9D?= =?UTF-8?q?=E8=B5=96=E6=B3=A8=E5=85=A5=E6=9E=B6=E6=9E=84?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 创建app包实现依赖注入容器和初始化器模式 - 重构main.go为六阶段清晰的初始化流程 - 新增结构化错误处理替代简陋的os.Exit调用 - 为HostInfo添加辅助函数增强功能但保持向后兼容 - 引入TargetInfo包装器支持上下文和元数据管理 - 优化代码组织提升可维护性和可测试性 --- Common/common.go | 52 +++++++++--------- Common/hostinfo_ext.go | 109 +++++++++++++++++++++++++++++++++++++ Common/target.go | 118 +++++++++++++++++++++++++++++++++++++++++ app/container.go | 107 +++++++++++++++++++++++++++++++++++++ app/errors.go | 43 +++++++++++++++ app/initializer.go | 66 +++++++++++++++++++++++ main.go | 92 +++++++++++++++++--------------- 7 files changed, 515 insertions(+), 72 deletions(-) create mode 100644 Common/hostinfo_ext.go create mode 100644 Common/target.go create mode 100644 app/container.go create mode 100644 app/errors.go create mode 100644 app/initializer.go diff --git a/Common/common.go b/Common/common.go index 7f42a59..63f4869 100644 --- a/Common/common.go +++ b/Common/common.go @@ -30,12 +30,8 @@ type ScanPlugin = base.ScanPlugin // 插件类型常量 const ( - PluginTypeService = base.PluginTypeService - PluginTypeWeb = base.PluginTypeWeb - PluginTypeLocal = base.PluginTypeLocal - PluginTypeBrute = base.PluginTypeBrute - PluginTypePoc = base.PluginTypePoc - PluginTypeScan = base.PluginTypeScan + PluginTypeWeb = base.PluginTypeWeb + PluginTypeLocal = base.PluginTypeLocal ) // 全局插件管理器 @@ -59,22 +55,22 @@ var loggerMutex sync.Mutex func getGlobalLogger() *logging.Logger { loggerMutex.Lock() defer loggerMutex.Unlock() - + if globalLogger == nil { level := getLogLevelFromString(LogLevel) config := &logging.LoggerConfig{ - Level: level, - EnableColor: !NoColor, - SlowOutput: false, - ShowProgress: ShowProgress, - StartTime: StartTime, + Level: level, + EnableColor: !NoColor, + SlowOutput: false, + ShowProgress: ShowProgress, + StartTime: StartTime, } globalLogger = logging.NewLogger(config) if ProgressBar != nil { globalLogger.SetProgressBar(ProgressBar) } globalLogger.SetOutputMutex(&OutputMutex) - + // 设置协调输出函数，使用LogWithProgress globalLogger.SetCoordinatedOutput(LogWithProgress) } @@ -105,18 +101,18 @@ func getLogLevelFromString(levelStr string) logging.LogLevel { } // 日志函数 -func InitLogger() { +func InitLogger() { loggerMutex.Lock() globalLogger = nil loggerMutex.Unlock() - getGlobalLogger().Initialize() + getGlobalLogger().Initialize() } -func LogDebug(msg string) { getGlobalLogger().Debug(msg) } -func LogBase(msg string) { getGlobalLogger().Base(msg) } -func LogInfo(msg string) { getGlobalLogger().Info(msg) } -func LogSuccess(result string) { getGlobalLogger().Success(result) } -func LogError(errMsg string) { getGlobalLogger().Error(errMsg) } +func LogDebug(msg string) { getGlobalLogger().Debug(msg) } +func LogBase(msg string) { getGlobalLogger().Base(msg) } +func LogInfo(msg string) { getGlobalLogger().Info(msg) } +func LogSuccess(result string) { getGlobalLogger().Success(result) } +func LogError(errMsg string) { getGlobalLogger().Error(errMsg) } // ============================================================================= // 输出系统简化接口 @@ -128,7 +124,7 @@ func InitOutput() error { if Outputfile == "" { return fmt.Errorf("output file not specified") } - + var format output.OutputFormat switch OutputFormat { case "txt": @@ -182,7 +178,7 @@ func WrapperTcpWithContext(ctx context.Context, network, address string) (net.Co Address: Socks5Proxy, Timeout: time.Second * 10, } - + proxyManager := proxy.NewProxyManager(proxyConfig) dialer, err := proxyManager.GetDialer() if err != nil { @@ -191,11 +187,11 @@ func WrapperTcpWithContext(ctx context.Context, network, address string) (net.Co var d net.Dialer return d.DialContext(ctx, network, address) } - + LogDebug(fmt.Sprintf("使用SOCKS5代理连接: %s -> %s", Socks5Proxy, address)) return dialer.DialContext(ctx, network, address) } - + // 没有配置代理，使用直连 var d net.Dialer return d.DialContext(ctx, network, address) @@ -210,7 +206,7 @@ func WrapperTlsWithContext(ctx context.Context, network, address string, config Address: Socks5Proxy, Timeout: time.Second * 10, } - + proxyManager := proxy.NewProxyManager(proxyConfig) tlsDialer, err := proxyManager.GetTLSDialer() if err != nil { @@ -219,11 +215,11 @@ func WrapperTlsWithContext(ctx context.Context, network, address string, config d := &tls.Dialer{Config: config} return d.DialContext(ctx, network, address) } - + LogDebug(fmt.Sprintf("使用SOCKS5代理TLS连接: %s -> %s", Socks5Proxy, address)) return tlsDialer.DialTLSContext(ctx, network, address, config) } - + // 没有配置代理，使用直连 d := &tls.Dialer{Config: config} return d.DialContext(ctx, network, address) @@ -236,4 +232,4 @@ func WrapperTlsWithContext(ctx context.Context, network, address string, config // CheckErrs 检查单个错误 - 简化版本 func CheckErrs(err error) error { return err -} \ No newline at end of file +} diff --git a/Common/hostinfo_ext.go b/Common/hostinfo_ext.go new file mode 100644 index 0000000..afab63c --- /dev/null +++ b/Common/hostinfo_ext.go @@ -0,0 +1,109 @@ +package common + +import ( + "fmt" + "strconv" +) + +// HostInfoHelper 提供HostInfo的辅助方法 +// 使用函数而不是方法，保持向后兼容 + +// GetHost 获取主机地址 +func GetHost(h *HostInfo) string { + return h.Host +} + +// GetPort 获取端口号（转换为整数） +func GetPort(h *HostInfo) (int, error) { + if h.Ports == "" { + return 0, fmt.Errorf("端口未设置") + } + return strconv.Atoi(h.Ports) +} + +// GetURL 获取URL地址 +func GetURL(h *HostInfo) string { + return h.Url +} + +// IsWebTarget 判断是否为Web目标 +func IsWebTarget(h *HostInfo) bool { + return h.Url != "" +} + +// HasPort 检查是否设置了端口 +func HasPort(h *HostInfo) bool { + return h.Ports != "" +} + +// HasHost 检查是否设置了主机 +func HasHost(h *HostInfo) bool { + return h.Host != "" +} + +// CloneHostInfo 克隆HostInfo（深拷贝） +func CloneHostInfo(h *HostInfo) HostInfo { + cloned := HostInfo{ + Host: h.Host, + Ports: h.Ports, + Url: h.Url, + } + + // 深拷贝Infostr切片 + if h.Infostr != nil { + cloned.Infostr = make([]string, len(h.Infostr)) + copy(cloned.Infostr, h.Infostr) + } + + return cloned +} + +// ValidateHostInfo 验证HostInfo的有效性 +func ValidateHostInfo(h *HostInfo) error { + if h.Host == "" && h.Url == "" { + return fmt.Errorf("主机地址或URL必须至少指定一个") + } + + // 验证端口格式（如果指定了） + if h.Ports != "" { + if _, err := GetPort(h); err != nil { + return fmt.Errorf("端口格式无效: %v", err) + } + } + + return nil +} + +// HostInfoString 返回HostInfo的字符串表示 +func HostInfoString(h *HostInfo) string { + if IsWebTarget(h) { + return h.Url + } + + if HasPort(h) { + return fmt.Sprintf("%s:%s", h.Host, h.Ports) + } + + return h.Host +} + +// AddInfo 添加附加信息 +func AddInfo(h *HostInfo, info string) { + if h.Infostr == nil { + h.Infostr = make([]string, 0) + } + h.Infostr = append(h.Infostr, info) +} + +// GetInfo 获取所有附加信息 +func GetInfo(h *HostInfo) []string { + if h.Infostr == nil { + return []string{} + } + return h.Infostr +} + +// HasInfo 检查是否有附加信息 +func HasInfo(h *HostInfo) bool { + return len(h.Infostr) > 0 +} \ No newline at end of file diff --git a/Common/target.go b/Common/target.go new file mode 100644 index 0000000..c0a18f2 --- /dev/null +++ b/Common/target.go @@ -0,0 +1,118 @@ +package common + +import ( + "context" +) + +// TargetInfo 包装HostInfo，提供更丰富的功能 +type TargetInfo struct { + *HostInfo // 嵌入HostInfo，保持向后兼容 + context context.Context + metadata map[string]interface{} +} + +// NewTargetInfo 创建新的目标信息 +func NewTargetInfo(hostInfo HostInfo) *TargetInfo { + return &TargetInfo{ + HostInfo: &hostInfo, + context: context.Background(), + metadata: make(map[string]interface{}), + } +} + +// NewTargetInfoFromPtr 从HostInfo指针创建目标信息 +func NewTargetInfoFromPtr(hostInfo *HostInfo) *TargetInfo { + return &TargetInfo{ + HostInfo: hostInfo, + context: context.Background(), + metadata: make(map[string]interface{}), + } +} + +// WithContext 设置上下文 +func (t *TargetInfo) WithContext(ctx context.Context) *TargetInfo { + t.context = ctx + return t +} + +// GetContext 获取上下文 +func (t *TargetInfo) GetContext() context.Context { + if t.context == nil { + t.context = context.Background() + } + return t.context +} + +// SetMetadata 设置元数据 +func (t *TargetInfo) SetMetadata(key string, value interface{}) *TargetInfo { + if t.metadata == nil { + t.metadata = make(map[string]interface{}) + } + t.metadata[key] = value + return t +} + +// GetMetadata 获取元数据 +func (t *TargetInfo) GetMetadata(key string) (interface{}, bool) { + if t.metadata == nil { + return nil, false + } + value, exists := t.metadata[key] + return value, exists +} + +// GetAllMetadata 获取所有元数据 +func (t *TargetInfo) GetAllMetadata() map[string]interface{} { + if t.metadata == nil { + return make(map[string]interface{}) + } + // 返回副本，防止外部修改 + result := make(map[string]interface{}) + for k, v := range t.metadata { + result[k] = v + } + return result +} + +// Clone 克隆目标信息 +func (t *TargetInfo) Clone() *TargetInfo { + clonedHost := CloneHostInfo(t.HostInfo) + cloned := &TargetInfo{ + HostInfo: &clonedHost, + context: t.context, + metadata: make(map[string]interface{}), + } + + // 复制元数据 + for k, v := range t.metadata { + cloned.metadata[k] = v + } + + return cloned +} + +// GetHostInfo 获取原始HostInfo（向后兼容） +func (t *TargetInfo) GetHostInfo() HostInfo { + return *t.HostInfo +} + +// Validate 验证目标信息 +func (t *TargetInfo) Validate() error { + return ValidateHostInfo(t.HostInfo) +} + +// String 返回字符串表示 +func (t *TargetInfo) String() string { + return HostInfoString(t.HostInfo) +} + +// IsValid 检查目标是否有效 +func (t *TargetInfo) IsValid() bool { + return t.Validate() == nil +} + +// HasMetadata 检查是否有指定的元数据 +func (t *TargetInfo) HasMetadata(key string) bool { + _, exists := t.GetMetadata(key) + return exists +} \ No newline at end of file diff --git a/app/container.go b/app/container.go new file mode 100644 index 0000000..710fb9d --- /dev/null +++ b/app/container.go @@ -0,0 +1,107 @@ +package app + +import ( + "context" + "fmt" + "sync" + + "github.com/shadow1ng/fscan/common" + "github.com/shadow1ng/fscan/core" +) + +// Container 依赖注入容器 +type Container struct { + services map[string]interface{} + initializers []Initializer + mu sync.RWMutex + initialized bool +} + +// NewContainer 创建新的容器 +func NewContainer() *Container { + container := &Container{ + services: make(map[string]interface{}), + } + + // 注册默认初始化器 + container.AddInitializer(&PluginInitializer{}) + + return container +} + +// AddInitializer 添加初始化器 +func (c *Container) AddInitializer(init Initializer) { + c.initializers = append(c.initializers, init) +} + +// Register 注册服务 +func (c *Container) Register(name string, service interface{}) { + c.mu.Lock() + defer c.mu.Unlock() + c.services[name] = service +} + +// Get 获取服务 +func (c *Container) Get(name string) (interface{}, bool) { + c.mu.RLock() + defer c.mu.RUnlock() + service, exists := c.services[name] + return service, exists +} + +// Initialize 初始化容器和所有服务 +func (c *Container) Initialize() error { + if c.initialized { + return nil + } + + // 执行所有初始化器 + for _, initializer := range c.initializers { + if err := initializer.Initialize(); err != nil { + return WrapError(ErrInitFailed, err) + } + } + + c.initialized = true + return nil +} + +// RunScan 执行扫描（包装现有的core.RunScan） +func (c *Container) RunScan(ctx context.Context, info common.HostInfo) error { + // 使用新的验证函数 + if err := common.ValidateHostInfo(&info); err != nil { + return WrapError(ErrScanFailed, err) + } + + // 创建目标信息（展示新功能，但保持兼容） + target := common.NewTargetInfo(info) + target.WithContext(ctx) + target.SetMetadata("container_managed", true) + target.SetMetadata("validation_passed", true) + + // 记录扫描信息 + c.logScanInfo(target) + + // 调用现有的扫描逻辑 + core.RunScan(info) + + return nil +} + +// logScanInfo 记录扫描信息 +func (c *Container) logScanInfo(target *common.TargetInfo) { + targetStr := target.String() + if targetStr != "" { + common.LogDebug(fmt.Sprintf("容器管理的扫描目标: %s", targetStr)) + } + + if target.HasMetadata("validation_passed") { + common.LogDebug("目标验证通过") + } +} + +// Cleanup 清理资源 +func (c *Container) Cleanup() { + // 清理输出资源 + common.CloseOutput() +} \ No newline at end of file diff --git a/app/errors.go b/app/errors.go new file mode 100644 index 0000000..af4d192 --- /dev/null +++ b/app/errors.go @@ -0,0 +1,43 @@ +package app + +import "fmt" + +// AppError 应用程序错误类型 +type AppError struct { + Code int + Message string + Cause error +} + +func (e *AppError) Error() string { + if e.Cause != nil { + return fmt.Sprintf("[%d] %s: %v", e.Code, e.Message, e.Cause) + } + return fmt.Sprintf("[%d] %s", e.Code, e.Message) +} + +// 预定义错误类型 +var ( + ErrInitFailed = &AppError{Code: 1, Message: "初始化失败"} + ErrParseFailed = &AppError{Code: 2, Message: "参数解析失败"} + ErrOutputFailed = &AppError{Code: 3, Message: "输出初始化失败"} + ErrScanFailed = &AppError{Code: 4, Message: "扫描执行失败"} +) + +// NewAppError 创建新的应用程序错误 +func NewAppError(code int, message string, cause error) *AppError { + return &AppError{ + Code: code, + Message: message, + Cause: cause, + } +} + +// WrapError 包装错误为应用程序错误 +func WrapError(baseErr *AppError, cause error) *AppError { + return &AppError{ + Code: baseErr.Code, + Message: baseErr.Message, + Cause: cause, + } +} \ No newline at end of file diff --git a/app/initializer.go b/app/initializer.go new file mode 100644 index 0000000..b62e19e --- /dev/null +++ b/app/initializer.go @@ -0,0 +1,66 @@ +package app + +import ( + "sort" + + "github.com/shadow1ng/fscan/common" + "github.com/shadow1ng/fscan/plugins/base" +) + +// Initializer 初始化器接口 +type Initializer interface { + Initialize() error + Name() string +} + +// PluginInitializer 插件初始化器 +type PluginInitializer struct{} + +func (p *PluginInitializer) Name() string { + return "PluginInitializer" +} + +func (p *PluginInitializer) Initialize() error { + var localPlugins []string + + // 获取所有注册的插件 + allPlugins := base.GlobalPluginRegistry.GetAll() + + for _, pluginName := range allPlugins { + metadata := base.GlobalPluginRegistry.GetMetadata(pluginName) + if metadata != nil && metadata.Category == "local" { + localPlugins = append(localPlugins, pluginName) + } + } + + // 排序以保持一致性 + sort.Strings(localPlugins) + + // 设置全局变量 + common.LocalPluginsList = localPlugins + + return nil +} + +// LoggerInitializer 日志初始化器 +type LoggerInitializer struct{} + +func (l *LoggerInitializer) Name() string { + return "LoggerInitializer" +} + +func (l *LoggerInitializer) Initialize() error { + common.InitLogger() + return nil +} + +// OutputInitializer 输出初始化器 +type OutputInitializer struct{} + +func (o *OutputInitializer) Name() string { + return "OutputInitializer" +} + +func (o *OutputInitializer) Initialize() error { + return common.InitOutput() +} \ No newline at end of file diff --git a/main.go b/main.go index 86330a9..32c16bf 100644 --- a/main.go +++ b/main.go @@ -1,58 +1,62 @@ package main import ( + "context" "fmt" "os" - "sort" + "github.com/shadow1ng/fscan/app" "github.com/shadow1ng/fscan/common" - "github.com/shadow1ng/fscan/core" - "github.com/shadow1ng/fscan/plugins/base" ) -// initLocalPlugins 初始化本地插件列表 -func initLocalPlugins() { - var localPlugins []string - - // 获取所有注册的插件 - allPlugins := base.GlobalPluginRegistry.GetAll() - - for _, pluginName := range allPlugins { - metadata := base.GlobalPluginRegistry.GetMetadata(pluginName) - if metadata != nil && metadata.Category == "local" { - localPlugins = append(localPlugins, pluginName) - } - } - - // 排序以保持一致性 - sort.Strings(localPlugins) - - // 设置全局变量 - common.LocalPluginsList = localPlugins -} - func main() { - // 初始化本地插件列表 - initLocalPlugins() + // 创建应用容器 + container := app.NewContainer() - var Info common.HostInfo - common.Flag(&Info) - - // 在flag解析后初始化logger，确保LogLevel参数生效 - common.InitLogger() - - // 解析 CLI 参数 - if err := common.Parse(&Info); err != nil { - os.Exit(1) + // 第一阶段：基础初始化（插件系统） + if err := container.Initialize(); err != nil { + handleError("基础初始化失败", err) } - - // 初始化输出系统，如果失败则直接退出 - if err := common.InitOutput(); err != nil { - common.LogError(fmt.Sprintf("初始化输出系统失败: %v", err)) - os.Exit(1) + defer container.Cleanup() + + // 第二阶段：解析配置 + var info common.HostInfo + common.Flag(&info) + + // 第三阶段：日志初始化（依赖于flag解析） + logInit := &app.LoggerInitializer{} + if err := logInit.Initialize(); err != nil { + handleError("日志初始化失败", err) + } + + // 第四阶段：参数解析和验证 + if err := common.Parse(&info); err != nil { + handleError("参数解析失败", err) + } + + // 第五阶段：输出系统初始化 + outputInit := &app.OutputInitializer{} + if err := outputInit.Initialize(); err != nil { + handleError("输出初始化失败", err) + } + + // 第六阶段：执行扫描 + ctx := context.Background() + if err := container.RunScan(ctx, info); err != nil { + handleError("扫描失败", err) + } +} + +func handleError(msg string, err error) { + // 检查是否是应用程序错误 + if appErr, ok := err.(*app.AppError); ok { + common.LogError(fmt.Sprintf("%s: %s", msg, appErr.Message)) + if appErr.Cause != nil { + common.LogError(fmt.Sprintf("详细错误: %v", appErr.Cause)) + } + os.Exit(appErr.Code) + } else { + common.LogError(fmt.Sprintf("%s: %v", msg, err)) + os.Exit(1) } - defer common.CloseOutput() - - // 执行 CLI 扫描逻辑 - core.RunScan(Info) }