diff --git a/src/codegate/cli.py b/src/codegate/cli.py index b230eba9..ea871b5e 100644 --- a/src/codegate/cli.py +++ b/src/codegate/cli.py @@ -305,16 +305,16 @@ def serve( ca = CertificateAuthority.get_instance() ca.ensure_certificates_exist() - # Set up event loop - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - # Initialize secrets manager and pipeline factory secrets_manager = SecretsManager() pipeline_factory = PipelineFactory(secrets_manager) app = init_app(pipeline_factory) + # Set up event loop + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + # Run the server try: loop.run_until_complete(run_servers(cfg, app)) diff --git a/src/codegate/pipeline/secrets/secrets.py b/src/codegate/pipeline/secrets/secrets.py index 67bc0386..0f3d61c2 100644 --- a/src/codegate/pipeline/secrets/secrets.py +++ b/src/codegate/pipeline/secrets/secrets.py @@ -26,6 +26,8 @@ class CodegateSecrets(PipelineStep): def __init__(self): """Initialize the CodegateSecrets pipeline step.""" super().__init__() + # Initialize and load signatures immediately + CodegateSignatures.initialize("signatures.yaml") @property def name(self) -> str: diff --git a/src/codegate/pipeline/secrets/signatures.py b/src/codegate/pipeline/secrets/signatures.py index 1f1a7ddd..95d306e1 100644 --- a/src/codegate/pipeline/secrets/signatures.py +++ b/src/codegate/pipeline/secrets/signatures.py @@ -39,7 +39,6 @@ class CodegateSignatures: """Main class for detecting secrets in text using regex patterns.""" _instance_lock: ClassVar[Lock] = Lock() - _signatures_loaded: ClassVar[bool] = False _signature_groups: ClassVar[List[SignatureGroup]] = [] _compiled_regexes: ClassVar[Dict[str, re.Pattern]] = {} _yaml_path: ClassVar[Optional[str]] = None @@ -48,7 +47,6 @@ class CodegateSignatures: def reset(cls) -> None: """Reset the cached patterns.""" with cls._instance_lock: - cls._signatures_loaded = False cls._signature_groups = [] cls._compiled_regexes = {} cls._yaml_path = None @@ -56,14 +54,16 @@ def reset(cls) -> None: @classmethod def initialize(cls, yaml_path: str) -> None: - """Initialize the SecretFinder with a YAML file path.""" + """Initialize the SecretFinder with a YAML file path and load signatures.""" if not Path(yaml_path).exists(): raise FileNotFoundError(f"Signatures file not found: {yaml_path}") with cls._instance_lock: - cls._yaml_path = yaml_path - cls._signatures_loaded = False - logger.debug(f"SecretFinder initialized with {yaml_path}") + # Only initialize if not already initialized with this path + if cls._yaml_path != yaml_path: + cls._yaml_path = yaml_path + cls._load_signatures() + logger.debug(f"SecretFinder initialized with {yaml_path}") @classmethod def _preprocess_yaml(cls, content: str) -> str: @@ -172,6 +172,10 @@ def _add_signature_group(cls, name: str, patterns: Dict[str, str]) -> None: def _load_signatures(cls) -> None: """Load signature patterns from the YAML file.""" try: + # Clear existing signatures before loading new ones + cls._signature_groups = [] + cls._compiled_regexes = {} + yaml_data = cls._load_yaml(cls._yaml_path) # Add custom GitHub token patterns @@ -205,32 +209,14 @@ def _load_signatures(cls) -> None: logger.error(f"Error loading signatures: {e}") raise - @classmethod - def _ensure_signatures_loaded(cls) -> None: - """Ensure signatures are loaded before use.""" - if not cls._signatures_loaded: - with cls._instance_lock: - if not cls._signatures_loaded: - if not cls._yaml_path: - raise RuntimeError("SecretFinder not initialized. Call initialize() first.") - try: - cls._load_signatures() - cls._signatures_loaded = True - except Exception as e: - logger.error(f"Failed to load signatures: {e}") - raise - @classmethod def find_in_string(cls, text: str) -> List[Match]: """Search for secrets in the provided string.""" if not text: return [] - try: - cls._ensure_signatures_loaded() - except Exception as e: - logger.error(f"Failed to load signatures: {e}") - return [] + if not cls._yaml_path: + raise RuntimeError("SecretFinder not initialized.") matches = [] lines = text.splitlines() diff --git a/src/codegate/server.py b/src/codegate/server.py index 34da0b8c..d2d7801e 100644 --- a/src/codegate/server.py +++ b/src/codegate/server.py @@ -4,7 +4,6 @@ from codegate import __description__, __version__ from codegate.dashboard.dashboard import dashboard_router from codegate.pipeline.factory import PipelineFactory -from codegate.pipeline.secrets.signatures import CodegateSignatures from codegate.providers.anthropic.provider import AnthropicProvider from codegate.providers.llamacpp.provider import LlamaCppProvider from codegate.providers.ollama.provider import OllamaProvider @@ -31,9 +30,6 @@ def init_app(pipeline_factory: PipelineFactory) -> FastAPI: # Create provider registry registry = ProviderRegistry(app) - # Initialize SignaturesFinder - CodegateSignatures.initialize("signatures.yaml") - # Register all known providers registry.add_provider( "openai", diff --git a/tests/pipeline/secrets/test_signatures.py b/tests/pipeline/secrets/test_signatures.py index c7e0b9e6..42f6470f 100644 --- a/tests/pipeline/secrets/test_signatures.py +++ b/tests/pipeline/secrets/test_signatures.py @@ -69,14 +69,12 @@ def test_initialize_and_reset(temp_yaml_file): CodegateSignatures.reset() assert CodegateSignatures._yaml_path is None - assert not CodegateSignatures._signatures_loaded assert not CodegateSignatures._signature_groups assert not CodegateSignatures._compiled_regexes def test_find_in_string_with_aws_credentials(temp_yaml_file): CodegateSignatures.initialize(temp_yaml_file) - CodegateSignatures._signatures_loaded = False # Force reload of signatures test_string = """ aws_access_key = 'AKIAIOSFODNN7EXAMPLE' @@ -100,7 +98,6 @@ def test_find_in_string_with_github_token(): try: CodegateSignatures.initialize(f.name) - CodegateSignatures._signatures_loaded = False # Force reload of signatures test_string = "github_token = 'ghp_1234567890abcdef1234567890abcdef123456'" matches = CodegateSignatures.find_in_string(test_string) @@ -121,7 +118,6 @@ def test_find_in_string_with_github_token(): def test_find_in_string_with_no_matches(temp_yaml_file): CodegateSignatures.initialize(temp_yaml_file) - CodegateSignatures._signatures_loaded = False # Force reload of signatures test_string = "No secrets here!" matches = CodegateSignatures.find_in_string(test_string) @@ -158,7 +154,6 @@ def test_duplicate_patterns(): try: CodegateSignatures.initialize(f.name) - CodegateSignatures._signatures_loaded = False # Force reload of signatures test_string = "aws_key = 'AKIAIOSFODNN7EXAMPLE'" matches = CodegateSignatures.find_in_string(test_string) diff --git a/tests/test_cli.py b/tests/test_cli.py index 8557a63c..caa3a215 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -105,13 +105,6 @@ def test_provider_registration(mock_registry, mock_secrets_mgr, mock_pipeline_fa assert "ollama" in provider_names -@patch("codegate.server.CodegateSignatures") -def test_signatures_initialization(mock_signatures, mock_pipeline_factory) -> None: - """Test that signatures are initialized correctly.""" - init_app(mock_pipeline_factory) - mock_signatures.initialize.assert_called_once_with("signatures.yaml") - - def test_pipeline_initialization(mock_pipeline_factory) -> None: """Test that pipelines are initialized correctly.""" app = init_app(mock_pipeline_factory) diff --git a/tests/test_server.py b/tests/test_server.py index 9170f36e..5420ad70 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -99,13 +99,6 @@ def test_provider_registration(mock_registry, mock_secrets_mgr, mock_pipeline_fa assert "ollama" in provider_names -@patch("codegate.server.CodegateSignatures") -def test_signatures_initialization(mock_signatures, mock_pipeline_factory) -> None: - """Test that signatures are initialized correctly.""" - init_app(mock_pipeline_factory) - mock_signatures.initialize.assert_called_once_with("signatures.yaml") - - def test_pipeline_initialization(mock_pipeline_factory) -> None: """Test that pipelines are initialized correctly.""" app = init_app(mock_pipeline_factory)