diff --git a/admin/admin.py b/admin/admin.py index e028053..42b0cc1 100644 --- a/admin/admin.py +++ b/admin/admin.py @@ -109,19 +109,62 @@ def load_config(self, toml_file: str) -> admin_service_pb2.LoadConfigResponse: raise Exception(f"Server error: {error_msg}") return response + def get_config_toml(self, output_file: str): + request = admin_service_pb2.GetConfigRequest() + response = self.stub.GetConfigAdmin(request) + + toml_data = response.config_data + + if output_file: + with open(output_file, "wb") as f: + f.write(toml_data) + print(f"TOML config saved to {output_file}") + + return toml_data + + def toggle_filtering(self, worker_id: int, enabled: bool): + request = admin_service_pb2.ToggleFilteringRequest( + worker_id=worker_id, + enabled=enabled, + ) + response = self.stub.ToggleFiltering(request) + return response + def main(): parser = argparse.ArgumentParser() - parser.add_argument("--file", default="config.toml") + subparsers = parser.add_subparsers(dest="action") + + load_parser = subparsers.add_parser("load") + load_parser.add_argument("--file", default="config.toml") + + get_parser = subparsers.add_parser("get") + get_parser.add_argument("--save", "-s", default="policy.toml") + + toggle_parser = subparsers.add_parser("toggle") + toggle_parser.add_argument("--id", type=int, required=True) + toggle_parser.add_argument("--on", action="store_true", dest="enabled") + toggle_parser.add_argument("--off", action="store_false", dest="enabled") + toggle_parser.set_defaults(enabled=True) + args = parser.parse_args() + client = AdminClient() try: - client = AdminClient() - client.load_config(args.file) - print("Config loaded") + if args.action == "load": + client.load_config(args.file) + print("Config loaded") + elif args.action == "get": + toml_data = client.get_config_toml(args.save) + print("Config saved") + elif args.action == "toggle": + resp = client.toggle_filtering(args.id, args.enabled) + state = "enabled" if args.enabled else "disabled" + print(f"Filtering {state}: {resp.message}") + else: + parser.print_help() except Exception as e: - print(f"Error loading config: {e}") + print(f"Error: {e}") return 1 - return 0 if __name__ == "__main__": diff --git a/admin/admin_service.proto b/admin/admin_service.proto index 77b597e..2ec01ac 100644 --- a/admin/admin_service.proto +++ b/admin/admin_service.proto @@ -6,6 +6,8 @@ option go_package = "github.com/moevm/grpc_server/pkg/proto/admin_service"; service AdminService { rpc LoadConfig(LoadConfigRequest) returns (LoadConfigResponse) {} + rpc GetConfigAdmin(GetConfigRequest) returns (GetConfigResponse) {} + rpc ToggleFiltering(ToggleFilteringRequest) returns (ToggleFilteringResponse) {} } message LoadConfigRequest { @@ -16,3 +18,19 @@ message LoadConfigResponse { bool success = 1; string error_message = 2; } + +message GetConfigRequest {} + +message GetConfigResponse { + bytes config_data = 1; +} + +message ToggleFilteringRequest { + uint64 worker_id = 1; + bool enabled = 2; +} + +message ToggleFilteringResponse { + bool success = 1; + string message = 2; +} diff --git a/controller/internal/grpcserver/admin_server.go b/controller/internal/grpcserver/admin_server.go index 262d2e2..f1cf751 100644 --- a/controller/internal/grpcserver/admin_server.go +++ b/controller/internal/grpcserver/admin_server.go @@ -2,7 +2,7 @@ package grpcserver import ( "context" - + "fmt" "github.com/moevm/grpc_server/internal/manager" pb "github.com/moevm/grpc_server/pkg/proto/admin_service" ) @@ -41,3 +41,24 @@ func (s *AdminServer) LoadConfig(ctx context.Context, req *pb.LoadConfigRequest) func (s *AdminServer) GetConfig() []byte { return s.configData } + +func (s *AdminServer) GetConfigAdmin(ctx context.Context, req *pb.GetConfigRequest) (*pb.GetConfigResponse, error) { + return &pb.GetConfigResponse{ + ConfigData: s.configData, + }, nil +} + +func (s *AdminServer) ToggleFiltering(ctx context.Context, req *pb.ToggleFilteringRequest) (*pb.ToggleFilteringResponse, error) { + if s.manager == nil { + return &pb.ToggleFilteringResponse{Success: false, Message: "manager not initialized"}, nil + } + + s.manager.SetFilteringEnabled(req.WorkerId, req.Enabled) + + return &pb.ToggleFilteringResponse{ + Success: true, + Message: fmt.Sprintf("Filtering %s for worker %d", + map[bool]string{true: "enabled", false: "disabled"}[req.Enabled], + req.WorkerId), + }, nil +} diff --git a/controller/internal/grpcserver/data_server.go b/controller/internal/grpcserver/data_server.go index 3db8292..a777dc3 100644 --- a/controller/internal/grpcserver/data_server.go +++ b/controller/internal/grpcserver/data_server.go @@ -112,3 +112,12 @@ func (s *DataServer) SendStats(ctx context.Context, report *pb.StatsReport) (*em report.WorkerId, report.TotalBlocked, report.TotalAllowed) return &emptypb.Empty{}, nil } + +func (s *DataServer) ToggleFiltering(ctx context.Context, req *pb.ToggleFilteringRequest) (*pb.ToggleFilteringResponse, error) { + s.manager.SetFilteringEnabled(req.WorkerId, req.Enabled) + enabled := s.manager.IsFilteringEnabled(req.WorkerId) + return &pb.ToggleFilteringResponse{ + Success: true, + Enabled: enabled, + }, nil +} diff --git a/controller/internal/manager/manager.go b/controller/internal/manager/manager.go index 5c99619..b38b199 100644 --- a/controller/internal/manager/manager.go +++ b/controller/internal/manager/manager.go @@ -3,18 +3,22 @@ package manager import ( "fmt" "log" + "sync" communication "github.com/moevm/grpc_server/pkg/proto/communication" "google.golang.org/protobuf/proto" ) type Manager struct { - policyManager *PolicyManager + policyManager *PolicyManager + filteringEnabled map[uint64]bool + mu sync.RWMutex } func NewManager() (*Manager, error) { return &Manager{ - policyManager: NewPolicyManager(), + policyManager: NewPolicyManager(), + filteringEnabled: make(map[uint64]bool), }, nil } @@ -51,3 +55,28 @@ func (m *Manager) UpdateConfig(configData []byte) error { } return m.policyManager.UpdateConfig(configData) } + +func (m *Manager) SetFilteringEnabled(workerID uint64, enabled bool) { + m.mu.Lock() + defer m.mu.Unlock() + + if workerID == 0 { + for id := range m.filteringEnabled { + m.filteringEnabled[id] = enabled + } + } else { + m.filteringEnabled[workerID] = enabled + } + log.Printf("Filtering %s for worker %d", + map[bool]string{true: "enabled", false: "disabled"}[enabled], workerID) +} + +func (m *Manager) IsFilteringEnabled(workerID uint64) bool { + m.mu.RLock() + defer m.mu.RUnlock() + enabled, exists := m.filteringEnabled[workerID] + if !exists { + return false + } + return enabled +} diff --git a/controller/pkg/proto/admin_service/admin_service.proto b/controller/pkg/proto/admin_service/admin_service.proto index 77b597e..e42f3fd 100644 --- a/controller/pkg/proto/admin_service/admin_service.proto +++ b/controller/pkg/proto/admin_service/admin_service.proto @@ -6,6 +6,8 @@ option go_package = "github.com/moevm/grpc_server/pkg/proto/admin_service"; service AdminService { rpc LoadConfig(LoadConfigRequest) returns (LoadConfigResponse) {} + rpc GetConfigAdmin(GetConfigRequest) returns (GetConfigResponse) {} + rpc ToggleFiltering(ToggleFilteringRequest) returns (ToggleFilteringResponse) {} } message LoadConfigRequest { @@ -16,3 +18,19 @@ message LoadConfigResponse { bool success = 1; string error_message = 2; } + +message GetConfigRequest {} + +message GetConfigResponse { + bytes config_data = 1; +} + +message ToggleFilteringRequest { + uint64 worker_id = 1; + bool enabled = 2; +} + +message ToggleFilteringResponse { + bool success = 1; + string message = 2; +} \ No newline at end of file diff --git a/controller/pkg/proto/communication/communication.proto b/controller/pkg/proto/communication/communication.proto index 1fe6dc4..03bf982 100644 --- a/controller/pkg/proto/communication/communication.proto +++ b/controller/pkg/proto/communication/communication.proto @@ -7,6 +7,7 @@ service DataService { rpc GetPolicy(GetPolicyRequest) returns (GetPolicyResponse); rpc Classify(ClassifyRequest) returns (ClassifyResponse); rpc SendStats(StatsReport) returns (google.protobuf.Empty); + rpc ToggleFiltering(ToggleFilteringRequest) returns (ToggleFilteringResponse); } message GetPolicyRequest { @@ -60,3 +61,13 @@ message StatsReport { uint64 total_allowed = 4; repeated ResourceStats resources = 5; } + +message ToggleFilteringRequest { + uint64 worker_id = 1; + bool enabled = 2; +} + +message ToggleFilteringResponse { + bool success = 1; + bool enabled = 2; +}