PM-pinou
83 lines
2.2 KiB
Python
83 lines
2.2 KiB
Python
"""Configuration loader with Pydantic model, YAML file, and caching."""
|
|
|
|
import pathlib
|
|
|
|
import yaml
|
|
from pydantic import BaseModel
|
|
|
|
|
|
class Config(BaseModel):
|
|
"""Typed configuration for TianXuan."""
|
|
|
|
class Server(BaseModel):
|
|
host: str = '0.0.0.0'
|
|
port: int = 80
|
|
debug: bool = False
|
|
|
|
class Data(BaseModel):
|
|
schema_strict: bool = False
|
|
recursive: bool = False
|
|
|
|
class Clustering(BaseModel):
|
|
algorithm: str = 'hdbscan'
|
|
min_cluster_size: int = 5
|
|
random_state: int = 42
|
|
|
|
class LLM(BaseModel):
|
|
enabled: bool = False
|
|
base_url: str = ''
|
|
api_key: str = ''
|
|
model: str = 'gpt-4'
|
|
|
|
class Entity(BaseModel):
|
|
subnet_masks: list[int] = []
|
|
ip_columns: list[str] = ['src_ip', 'dst_ip']
|
|
|
|
entity: Entity = Entity()
|
|
server: Server = Server()
|
|
data: Data = Data()
|
|
clustering: Clustering = Clustering()
|
|
llm: LLM = LLM()
|
|
|
|
|
|
_config_path = pathlib.Path(__file__).parent / 'config.yaml'
|
|
_config_cache = None
|
|
_config_mtime = 0
|
|
|
|
|
|
def get_config() -> Config:
|
|
"""Return the cached Config, reloading if the YAML file has changed.
|
|
|
|
If *config.yaml* does not exist, returns an empty ``Config()`` with
|
|
all default values.
|
|
"""
|
|
global _config_cache, _config_mtime
|
|
|
|
if _config_path.exists():
|
|
mtime = _config_path.stat().st_mtime
|
|
if _config_cache is None or mtime > _config_mtime:
|
|
with open(_config_path, 'r', encoding='utf-8') as f:
|
|
data = yaml.safe_load(f) or {}
|
|
_config_cache = Config(**data)
|
|
_config_mtime = mtime
|
|
else:
|
|
if _config_cache is None:
|
|
_config_cache = Config()
|
|
|
|
return _config_cache
|
|
|
|
|
|
def save_config(cfg: Config) -> None:
|
|
"""Persist a Config instance back to *config.yaml* and update the cache."""
|
|
global _config_cache, _config_mtime
|
|
with open(_config_path, 'w', encoding='utf-8') as f:
|
|
yaml.dump(
|
|
cfg.model_dump(mode='python'),
|
|
f,
|
|
default_flow_style=False,
|
|
allow_unicode=True,
|
|
sort_keys=False,
|
|
)
|
|
_config_cache = cfg
|
|
_config_mtime = _config_path.stat().st_mtime
|