feat: 添加API地址和加密密钥配置,重构API服务启动逻辑

This commit is contained in:
tongque 2025-04-14 14:32:51 +08:00
parent f921d81a76
commit 16e40fe7ed
5 changed files with 68 additions and 31 deletions

View File

@ -957,6 +957,8 @@ var (
ShowScanPlan bool // 是否显示扫描计划详情 ShowScanPlan bool // 是否显示扫描计划详情
SlowLogOutput bool // 是否启用慢速日志输出 SlowLogOutput bool // 是否启用慢速日志输出
Language string // 界面语言设置 Language string // 界面语言设置
ApiAddr string // API地址
SecretKey string // 加密密钥
) )
var ( var (

View File

@ -152,6 +152,8 @@ func Flag(Info *HostInfo) {
// ═════════════════════════════════════════════════ // ═════════════════════════════════════════════════
flag.StringVar(&Shellcode, "sc", "", GetText("flag_shellcode")) flag.StringVar(&Shellcode, "sc", "", GetText("flag_shellcode"))
flag.StringVar(&Language, "lang", "zh", GetText("flag_language")) flag.StringVar(&Language, "lang", "zh", GetText("flag_language"))
flag.StringVar(&ApiAddr, "api", "", GetText("flag_api"))
flag.StringVar(&SecretKey, "secret", "", GetText("flag_api_key"))
// 解析命令行参数 // 解析命令行参数
parseCommandLineArgs() parseCommandLineArgs()

View File

@ -2,60 +2,67 @@ package rpc
import ( import (
"context" "context"
"log"
"net" "net"
"net/http" "net/http"
"github.com/grpc-ecosystem/grpc-gateway/v2/runtime" "github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
"github.com/shadow1ng/fscan/Common"
pb "github.com/shadow1ng/fscan/RPC/lib" pb "github.com/shadow1ng/fscan/RPC/lib"
"github.com/shadow1ng/fscan/RPC/service" "github.com/shadow1ng/fscan/RPC/service"
"google.golang.org/grpc" "google.golang.org/grpc"
"google.golang.org/grpc/credentials/insecure" "google.golang.org/grpc/credentials/insecure"
) )
// 暴露的启动函数(可供外部调用) // 启动 gRPC + HTTP Gateway 服务(仅当设置了 API 地址时)
func StartHTTPServer() { // 如果未设置 API 地址,直接返回 nil
go runGRPCServer() // 启动 gRPC 服务 // 如果 HTTP 启动失败,则返回 error
if err := runHTTPGateway(); err != nil { func StartApiServer() error {
log.Fatalf("HTTP 启动失败: %v", err) 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 服务 // 启动 gRPC 服务
func runGRPCServer() { func runGRPCServer(addr string) {
lis, err := net.Listen("tcp", ":50051") lis, err := net.Listen("tcp", addr)
if err != nil { if err != nil {
log.Fatalf("监听失败: %v", err) Common.LogError("监听失败: " + err.Error())
return
} }
s := grpc.NewServer() s := grpc.NewServer()
pb.RegisterFscanServiceServer(s, &service.FscanService{}) pb.RegisterFscanServiceServer(s, &service.FscanService{})
log.Println("✅ gRPC 服务已启动,端口 50051") Common.LogSuccess("✅ gRPC 服务已启动,地址: " + addr)
if err := s.Serve(lis); err != nil { if err := s.Serve(lis); err != nil {
log.Fatalf("gRPC 启动失败: %v", err) Common.LogError("gRPC 启动失败: " + err.Error())
} }
} }
// 启动 HTTP Gateway 服务 // 启动 HTTP Gateway
func runHTTPGateway() error { func runHTTPGateway(httpAddr, grpcAddr string) error {
ctx := context.Background() ctx := context.Background()
mux := runtime.NewServeMux() mux := runtime.NewServeMux()
opts := []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())} 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 { if err != nil {
return err return err
} }
// 包裹 mux加上 CORS 支持
handler := allowCORS(mux)
log.Println("✅ HTTP Gateway 已启动,端口 8080")
return http.ListenAndServe(":8080", handler)
}
// 添加 CORS 支持 // 添加 CORS 支持
func allowCORS(h http.Handler) http.Handler { handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Access-Control-Allow-Origin", "*") 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-Methods", "GET, POST, PUT, DELETE, OPTIONS")
w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization") 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) w.WriteHeader(http.StatusOK)
return return
} }
mux.ServeHTTP(w, r)
h.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
} }

View File

@ -2,9 +2,9 @@ package service
import ( import (
"context" "context"
"log"
"time" "time"
"github.com/shadow1ng/fscan/Common"
pb "github.com/shadow1ng/fscan/RPC/lib" pb "github.com/shadow1ng/fscan/RPC/lib"
) )
@ -21,8 +21,7 @@ type FscanService struct {
// - StartScanResponse包含任务 ID 和提示信息。 // - StartScanResponse包含任务 ID 和提示信息。
// - error执行中出现的错误信息。 // - error执行中出现的错误信息。
func (s *FscanService) StartScan(ctx context.Context, req *pb.StartScanRequest) (*pb.StartScanResponse, 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: 在此处实现实际的扫描逻辑,例如调用扫描器、创建任务、存储任务状态等。 // TODO: 在此处实现实际的扫描逻辑,例如调用扫描器、创建任务、存储任务状态等。
// 可以异步执行扫描逻辑,并生成一个唯一的 taskID 进行标识。 // 可以异步执行扫描逻辑,并生成一个唯一的 taskID 进行标识。

15
main.go
View File

@ -2,9 +2,11 @@ package main
import ( import (
"fmt" "fmt"
"os"
"github.com/shadow1ng/fscan/Common" "github.com/shadow1ng/fscan/Common"
"github.com/shadow1ng/fscan/Core" "github.com/shadow1ng/fscan/Core"
"os" rpc "github.com/shadow1ng/fscan/RPC"
) )
func main() { func main() {
@ -12,14 +14,23 @@ func main() {
var Info Common.HostInfo var Info Common.HostInfo
Common.Flag(&Info) Common.Flag(&Info)
// 启动 gRPC + HTTP Gateway 服务
if err := rpc.StartApiServer(); err != nil {
os.Exit(1)
}
// 解析 CLI 参数
if err := Common.Parse(&Info); err != nil { if err := Common.Parse(&Info); err != nil {
os.Exit(1) os.Exit(1)
} }
// 初始化输出系统,如果失败则直接退出 // 初始化输出系统,如果失败则直接退出
if err := Common.InitOutput(); err != nil { if err := Common.InitOutput(); err != nil {
Common.LogError(fmt.Sprintf("初始化输出系统失败: %v", err)) Common.LogError(fmt.Sprintf("初始化输出系统失败: %v", err))
os.Exit(1) // 关键修改:初始化失败时直接退出 os.Exit(1)
} }
defer Common.CloseOutput() defer Common.CloseOutput()
// 执行 CLI 扫描逻辑
Core.Scan(Info) Core.Scan(Info)
} }