From a30cd12249ebcdd5252ef2601cd373a8db372d8e Mon Sep 17 00:00:00 2001 From: tongque <2863528786@qq.com> Date: Tue, 22 Apr 2025 12:53:06 +0800 Subject: [PATCH] =?UTF-8?q?refactor:=20rpc=E8=AF=B7=E6=B1=82=E9=9C=80?= =?UTF-8?q?=E8=A6=81=E9=85=8D=E7=BD=AE=E8=AF=B7=E6=B1=82=E5=A4=B4:?= =?UTF-8?q?=E7=A7=98=E9=92=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- RPC/lib/rpc.pb.go | 34 ++++++++-------------------------- RPC/lib/rpc.proto | 6 ++---- RPC/server.go | 16 +++++++++++++++- RPC/service/fscan.go | 2 +- 4 files changed, 26 insertions(+), 32 deletions(-) diff --git a/RPC/lib/rpc.pb.go b/RPC/lib/rpc.pb.go index ddc7bc1..d692b0d 100644 --- a/RPC/lib/rpc.pb.go +++ b/RPC/lib/rpc.pb.go @@ -26,8 +26,7 @@ const ( // 启动任务的请求 type StartScanRequest struct { state protoimpl.MessageState `protogen:"open.v1"` - Secret string `protobuf:"bytes,1,opt,name=secret,proto3" json:"secret,omitempty"` - Arg string `protobuf:"bytes,2,opt,name=arg,proto3" json:"arg,omitempty"` + Arg string `protobuf:"bytes,1,opt,name=arg,proto3" json:"arg,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -62,13 +61,6 @@ func (*StartScanRequest) Descriptor() ([]byte, []int) { return file_lib_rpc_proto_rawDescGZIP(), []int{0} } -func (x *StartScanRequest) GetSecret() string { - if x != nil { - return x.Secret - } - return "" -} - func (x *StartScanRequest) GetArg() string { if x != nil { return x.Arg @@ -132,8 +124,7 @@ func (x *StartScanResponse) GetMessage() string { // 获取扫描结果的请求 type TaskResultsRequest struct { state protoimpl.MessageState `protogen:"open.v1"` - Secret string `protobuf:"bytes,1,opt,name=secret,proto3" json:"secret,omitempty"` // 用于身份校验 - Filter *Filter `protobuf:"bytes,2,opt,name=filter,proto3" json:"filter,omitempty"` // 筛选条件（如关键字、状态等） + Filter *Filter `protobuf:"bytes,1,opt,name=filter,proto3" json:"filter,omitempty"` // 筛选条件（如关键字、状态等） unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -168,13 +159,6 @@ func (*TaskResultsRequest) Descriptor() ([]byte, []int) { return file_lib_rpc_proto_rawDescGZIP(), []int{2} } -func (x *TaskResultsRequest) GetSecret() string { - if x != nil { - return x.Secret - } - return "" -} - func (x *TaskResultsRequest) GetFilter() *Filter { if x != nil { return x.Filter @@ -384,16 +368,14 @@ var File_lib_rpc_proto protoreflect.FileDescriptor const file_lib_rpc_proto_rawDesc = "" + "\n" + - "\rlib/rpc.proto\x12\x03lib\x1a\x1cgoogle/api/annotations.proto\x1a\x1cgoogle/protobuf/struct.proto\"<\n" + - "\x10StartScanRequest\x12\x16\n" + - "\x06secret\x18\x01 \x01(\tR\x06secret\x12\x10\n" + - "\x03arg\x18\x02 \x01(\tR\x03arg\"F\n" + + "\rlib/rpc.proto\x12\x03lib\x1a\x1cgoogle/api/annotations.proto\x1a\x1cgoogle/protobuf/struct.proto\"$\n" + + "\x10StartScanRequest\x12\x10\n" + + "\x03arg\x18\x01 \x01(\tR\x03arg\"F\n" + "\x11StartScanResponse\x12\x17\n" + "\atask_id\x18\x01 \x01(\tR\x06taskId\x12\x18\n" + - "\amessage\x18\x02 \x01(\tR\amessage\"Q\n" + - "\x12TaskResultsRequest\x12\x16\n" + - "\x06secret\x18\x01 \x01(\tR\x06secret\x12#\n" + - "\x06filter\x18\x02 \x01(\v2\v.lib.FilterR\x06filter\"[\n" + + "\amessage\x18\x02 \x01(\tR\amessage\"9\n" + + "\x12TaskResultsRequest\x12#\n" + + "\x06filter\x18\x01 \x01(\v2\v.lib.FilterR\x06filter\"[\n" + "\x06Filter\x12\x17\n" + "\atask_id\x18\x01 \x01(\tR\x06taskId\x12\x1d\n" + "\n" + diff --git a/RPC/lib/rpc.proto b/RPC/lib/rpc.proto index cd8a81a..92cb44a 100644 --- a/RPC/lib/rpc.proto +++ b/RPC/lib/rpc.proto @@ -34,8 +34,7 @@ service FscanService { // 启动任务的请求 message StartScanRequest { - string secret = 1; - string arg= 2; + string arg= 1; } // 启动任务的响应 @@ -46,8 +45,7 @@ message StartScanResponse { // 获取扫描结果的请求 message TaskResultsRequest { - string secret = 1; // 用于身份校验 - Filter filter = 2; // 筛选条件（如关键字、状态等） + Filter filter = 1; // 筛选条件（如关键字、状态等） } message Filter { diff --git a/RPC/server.go b/RPC/server.go index 7d43d27..f6b43f3 100644 --- a/RPC/server.go +++ b/RPC/server.go @@ -4,6 +4,7 @@ import ( "context" "net" "net/http" + "time" "github.com/grpc-ecosystem/grpc-gateway/v2/runtime" "github.com/shadow1ng/fscan/Common" @@ -13,11 +14,16 @@ import ( "google.golang.org/grpc/credentials/insecure" ) +var internalSecretKey string + // 启动 gRPC + HTTP Gateway 服务（仅当设置了 API 地址时） func StartApiServer() error { if Common.ApiAddr == "" { return nil } + if Common.SecretKey == "" { + internalSecretKey = time.Now().Format("20060102150405") + } grpcAddr := "127.0.0.1:50051" httpAddr := validateHTTPAddr(Common.ApiAddr, ":8088") @@ -61,7 +67,8 @@ func runHTTPGateway(httpAddr, grpcAddr string) error { // 使用中间件包装 mux handler := applyMiddlewares(mux) - Common.LogSuccess("✅ HTTP Gateway 已启动，地址: " + httpAddr) + Common.LogSuccess("✅ HTTP Gateway 已启动，地址: http://" + httpAddr) + Common.LogSuccess("✅ API Secret: " + internalSecretKey) return http.ListenAndServe(httpAddr, handler) } @@ -77,6 +84,13 @@ func applyMiddlewares(handler http.Handler) http.Handler { return } + secret := r.Header.Get("Fscan-API-SECRET") + if secret == "" || secret != internalSecretKey { + http.Error(w, `无效的 API Secret，请通过请求头 Fscan-API-SECRET 提供正确的密钥。 + 如果你未手动配置 SecretKey，服务会在启动时自动生成一个随机密钥，并输出到日志中。`, http.StatusUnauthorized) + return + } + handler.ServeHTTP(w, r) }) } diff --git a/RPC/service/fscan.go b/RPC/service/fscan.go index 471b8a0..935e05d 100644 --- a/RPC/service/fscan.go +++ b/RPC/service/fscan.go @@ -36,7 +36,7 @@ func (s *FscanService) StartScan(ctx context.Context, req *pb.StartScanRequest) s.scanMutex.Lock() defer s.scanMutex.Unlock() - Common.LogDebug("异步执行扫描请求，目标: " + req.Arg + ", " + req.Secret) + Common.LogDebug("异步执行扫描请求，目标: " + req.Arg) var info Common.HostInfo if err := Common.FlagFromRemote(&info, req.Arg); err != nil {