From 3435d7e6bf7b9aed22c77d0c42b01354610cac97 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Thu, 20 Feb 2025 14:06:32 +0200
Subject: [PATCH] feat(simplify lib by refactoring config):

---
 agentic_security/config.py | 65 ++++++++++++++++++++++++++++++++++++++
 agentic_security/lib.py    | 65 +-------------------------------------
 2 files changed, 66 insertions(+), 64 deletions(-)
 create mode 100644 agentic_security/config.py

diff --git a/agentic_security/config.py b/agentic_security/config.py
new file mode 100644
index 0000000..f5d8341
--- /dev/null
+++ b/agentic_security/config.py
@@ -0,0 +1,65 @@
+import tomli
+from loguru import logger
+
+
+class CfgMixin:
+    config = {}
+    default_path = "agentic_security.toml"
+
+    def get_or_create_config(self) -> bool:
+        if not self.has_local_config():
+            self.generate_default_cfg()
+            return False
+        self.load_config(self.default_path)
+        return True
+
+    def has_local_config(self):
+        try:
+            with open(self.default_path):
+                return True
+        except FileNotFoundError:
+            return False
+
+    @classmethod
+    def load_config(cls, config_path: str):
+        """
+        Load configuration from a TOML file and store it in the class variable.
+
+        Args:
+            config_path (str): Path to the TOML configuration file.
+
+        Raises:
+            FileNotFoundError: If the configuration file is not found.
+            toml.TomlDecodeError: If the configuration file has syntax errors.
+        """
+        try:
+            with open(config_path, "rb") as config_file:
+                cls.config = tomli.load(config_file)
+                logger.info(f"Configuration loaded successfully from {config_path}.")
+        except FileNotFoundError:
+            logger.error(f"Configuration file {config_path} not found.")
+            raise
+        except Exception as e:
+            logger.error(f"Error parsing TOML configuration: {e}")
+            raise
+
+    @classmethod
+    def get_config_value(cls, key: str, default=None):
+        """
+        Retrieve a configuration value by key from the loaded configuration.
+
+        Args:
+            key (str): Dot-separated key path to the configuration value (e.g., 'general.maxBudget').
+            default: Default value if the key is not found.
+
+        Returns:
+            The configuration value if found, otherwise the default value.
+        """
+        keys = key.split(".")
+        value = cls.config
+        for k in keys:
+            if isinstance(value, dict) and k in value:
+                value = value[k]
+            else:
+                return default
+        return value
diff --git a/agentic_security/lib.py b/agentic_security/lib.py
index 154641c..a8f066f 100644
--- a/agentic_security/lib.py
+++ b/agentic_security/lib.py
@@ -3,7 +3,6 @@ import json
 from datetime import datetime
 
 import colorama
-import tomli
 import tqdm.asyncio
 from loguru import logger
 from rich.console import Console
@@ -13,6 +12,7 @@ from tabulate import tabulate
 from agentic_security.models.schemas import Scan
 from agentic_security.probe_data import REGISTRY
 from agentic_security.routes.scan import streaming_response_generator
+from agentic_security.config import CfgMixin  # Importing the configuration mixin
 
 # Enhanced color and style definitions
 RESET = colorama.Style.RESET_ALL
@@ -23,69 +23,6 @@ YELLOW = colorama.Fore.YELLOW
 BLUE = colorama.Fore.BLUE
 
 
-class CfgMixin:
-    config = {}
-    default_path = "agentic_security.toml"
-
-    def get_or_create_config(self) -> bool:
-        if not self.has_local_config():
-            self.generate_default_cfg()
-            return False
-        self.load_config(self.default_path)
-        return True
-
-    def has_local_config(self):
-        try:
-            with open(self.default_path):
-                return True
-        except FileNotFoundError:
-            return False
-
-    @classmethod
-    def load_config(cls, config_path: str):
-        """
-        Load configuration from a TOML file and store it in the class variable.
-
-        Args:
-            config_path (str): Path to the TOML configuration file.
-
-        Raises:
-            FileNotFoundError: If the configuration file is not found.
-            toml.TomlDecodeError: If the configuration file has syntax errors.
-        """
-        try:
-            with open(config_path, "rb") as config_file:
-                cls.config = tomli.load(config_file)
-                logger.info(f"Configuration loaded successfully from {config_path}.")
-        except FileNotFoundError:
-            logger.error(f"Configuration file {config_path} not found.")
-            raise
-        except Exception as e:
-            logger.error(f"Error parsing TOML configuration: {e}")
-            raise
-
-    @classmethod
-    def get_config_value(cls, key: str, default=None):
-        """
-        Retrieve a configuration value by key from the loaded configuration.
-
-        Args:
-            key (str): Dot-separated key path to the configuration value (e.g., 'general.maxBudget').
-            default: Default value if the key is not found.
-
-        Returns:
-            The configuration value if found, otherwise the default value.
-        """
-        keys = key.split(".")
-        value = cls.config
-        for k in keys:
-            if isinstance(value, dict) and k in value:
-                value = value[k]
-            else:
-                return default
-        return value
-
-
 class AgenticSecurity(CfgMixin):
     @classmethod
     async def async_scan(