diff --git a/Common/Config.go b/Common/Config.go index 921d95d..c9491b2 100644 --- a/Common/Config.go +++ b/Common/Config.go @@ -957,6 +957,8 @@ var ( ShowScanPlan bool // 是否显示扫描计划详情 SlowLogOutput bool // 是否启用慢速日志输出 Language string // 界面语言设置 + ApiAddr string // API地址 + SecretKey string // 加密密钥 ) var ( diff --git a/Common/Flag.go b/Common/Flag.go index 66af14a..f82abac 100644 --- a/Common/Flag.go +++ b/Common/Flag.go @@ -152,6 +152,8 @@ func Flag(Info *HostInfo) { // ═════════════════════════════════════════════════ flag.StringVar(&Shellcode, "sc", "", GetText("flag_shellcode")) flag.StringVar(&Language, "lang", "zh", GetText("flag_language")) + flag.StringVar(&ApiAddr, "api", "", GetText("flag_api")) + flag.StringVar(&SecretKey, "secret", "", GetText("flag_api_key")) // 解析命令行参数 parseCommandLineArgs() diff --git a/RPC/server.go b/RPC/server.go index 7fefb04..3555724 100644 --- a/RPC/server.go +++ b/RPC/server.go @@ -2,60 +2,67 @@ package rpc import ( "context" - "log" "net" "net/http" "github.com/grpc-ecosystem/grpc-gateway/v2/runtime" + "github.com/shadow1ng/fscan/Common" pb "github.com/shadow1ng/fscan/RPC/lib" "github.com/shadow1ng/fscan/RPC/service" "google.golang.org/grpc" "google.golang.org/grpc/credentials/insecure" ) -// 暴露的启动函数(可供外部调用) -func StartHTTPServer() { - go runGRPCServer() // 启动 gRPC 服务 - if err := runHTTPGateway(); err != nil { - log.Fatalf("HTTP 启动失败: %v", err) +// 启动 gRPC + HTTP Gateway 服务(仅当设置了 API 地址时) +// 如果未设置 API 地址,直接返回 nil +// 如果 HTTP 启动失败,则返回 error +func StartApiServer() error { + if Common.ApiAddr == "" { + Common.LogDebug("未设置 API 地址,跳过 API 服务启动") + return nil } + + grpcAddr := ":50051" + httpAddr := validateHTTPAddr(Common.ApiAddr, ":8088") + + go runGRPCServer(grpcAddr) + + if err := runHTTPGateway(httpAddr, grpcAddr); err != nil { + Common.LogError("HTTP 启动失败: " + err.Error()) + return err + } + + return nil } // 启动 gRPC 服务 -func runGRPCServer() { - lis, err := net.Listen("tcp", ":50051") +func runGRPCServer(addr string) { + lis, err := net.Listen("tcp", addr) if err != nil { - log.Fatalf("监听失败: %v", err) + Common.LogError("监听失败: " + err.Error()) + return } s := grpc.NewServer() pb.RegisterFscanServiceServer(s, &service.FscanService{}) - log.Println("✅ gRPC 服务已启动,端口 50051") + Common.LogSuccess("✅ gRPC 服务已启动,地址: " + addr) if err := s.Serve(lis); err != nil { - log.Fatalf("gRPC 启动失败: %v", err) + Common.LogError("gRPC 启动失败: " + err.Error()) } } -// 启动 HTTP Gateway 服务 -func runHTTPGateway() error { +// 启动 HTTP Gateway +func runHTTPGateway(httpAddr, grpcAddr string) error { ctx := context.Background() mux := runtime.NewServeMux() opts := []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())} - err := pb.RegisterFscanServiceHandlerFromEndpoint(ctx, mux, "localhost:50051", opts) + err := pb.RegisterFscanServiceHandlerFromEndpoint(ctx, mux, grpcAddr, opts) if err != nil { return err } - // 包裹 mux,加上 CORS 支持 - handler := allowCORS(mux) - - log.Println("✅ HTTP Gateway 已启动,端口 8080") - return http.ListenAndServe(":8080", handler) -} - -// 添加 CORS 支持 -func allowCORS(h http.Handler) http.Handler { - return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // 添加 CORS 支持 + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Access-Control-Allow-Origin", "*") w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS") w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization") @@ -64,7 +71,23 @@ func allowCORS(h http.Handler) http.Handler { w.WriteHeader(http.StatusOK) return } - - h.ServeHTTP(w, r) + mux.ServeHTTP(w, r) }) + + Common.LogSuccess("✅ HTTP Gateway 已启动,地址: " + httpAddr) + return http.ListenAndServe(httpAddr, handler) +} + +// 校验监听地址格式,格式非法使用默认 +func validateHTTPAddr(input, fallback string) string { + if input == "" { + Common.LogInfo("未指定 API 地址,使用默认地址: " + fallback) + return fallback + } + _, _, err := net.SplitHostPort(input) + if err != nil { + Common.LogError("无效的 API 地址格式 [" + input + "],使用默认地址: " + fallback) + return fallback + } + return input } diff --git a/RPC/service/fscan.go b/RPC/service/fscan.go index 74e2e33..5d2aa29 100644 --- a/RPC/service/fscan.go +++ b/RPC/service/fscan.go @@ -2,9 +2,9 @@ package service import ( "context" - "log" "time" + "github.com/shadow1ng/fscan/Common" pb "github.com/shadow1ng/fscan/RPC/lib" ) @@ -21,8 +21,7 @@ type FscanService struct { // - StartScanResponse:包含任务 ID 和提示信息。 // - error:执行中出现的错误信息。 func (s *FscanService) StartScan(ctx context.Context, req *pb.StartScanRequest) (*pb.StartScanResponse, error) { - log.Printf("启动扫描: target=%s, port=%s, mode=%s", req.Secret, req.Arg) - + Common.LogDebug("接收到扫描请求,目标: " + req.Arg + ", " + req.Secret) // TODO: 在此处实现实际的扫描逻辑,例如调用扫描器、创建任务、存储任务状态等。 // 可以异步执行扫描逻辑,并生成一个唯一的 taskID 进行标识。 diff --git a/main.go b/main.go index dea1f10..da2d8e2 100644 --- a/main.go +++ b/main.go @@ -2,9 +2,11 @@ package main import ( "fmt" + "os" + "github.com/shadow1ng/fscan/Common" "github.com/shadow1ng/fscan/Core" - "os" + rpc "github.com/shadow1ng/fscan/RPC" ) func main() { @@ -12,14 +14,23 @@ func main() { var Info Common.HostInfo Common.Flag(&Info) + + // 启动 gRPC + HTTP Gateway 服务 + if err := rpc.StartApiServer(); err != nil { + os.Exit(1) + } + // 解析 CLI 参数 if err := Common.Parse(&Info); err != nil { os.Exit(1) } + // 初始化输出系统,如果失败则直接退出 if err := Common.InitOutput(); err != nil { Common.LogError(fmt.Sprintf("初始化输出系统失败: %v", err)) - os.Exit(1) // 关键修改:初始化失败时直接退出 + os.Exit(1) } defer Common.CloseOutput() + + // 执行 CLI 扫描逻辑 Core.Scan(Info) }