feat: update api in controller

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__":

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;
+}

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
+}

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
+}

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
+}

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;
+}

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;
+}