feat(simplify lib by refactoring config):

This commit is contained in:
Alexander Myasoedov
2025-02-20 14:06:32 +02:00
parent ee3faab415
commit 3435d7e6bf
2 changed files with 66 additions and 64 deletions
+65
View File
@@ -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
+1 -64
View File
@@ -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(