import logging from defence360agent.contracts.config import IConfig, LocalConfig logger = logging.getLogger(__name__) def _fix_config(config_file): try: config = config_file.config_to_dict(normalize=False) if "DOS" not in config: return config["DOS"].pop("timeout", None) config["DOS"].pop("max_connections", None) config_file.dict_to_config(config, overwrite=True, validate=False) except Exception: logger.exception("Failed to remove fields") def migrate( migrator, database, fake=False, config_file: IConfig = LocalConfig(), **kwargs ): if fake: return _fix_config(config_file) def rollback(migrator, database, fake=False, **kwargs): pass