import re from typing import Optional from defence360agent.utils import atomic_rewrite class KWConfig: """ Basic class for working with key-value configuration files Subclasses must define SEARCH_PATTERN and WRITE_PATTERN attributes """ SEARCH_PATTERN = DEFAULT_FILENAME = WRITE_PATTERN = "" ALLOW_EMPTY_CONFIG = True def __init__(self, name, filename=None): assert self.SEARCH_PATTERN self._pattern = re.compile( self.SEARCH_PATTERN.format(name), re.MULTILINE ) self._filename = filename or self.DEFAULT_FILENAME self._name = name def set(self, value) -> Optional[str]: assert self.WRITE_PATTERN with open(self._filename) as f: content = f.read() old_value = self._parse(content) if old_value is None: # If no variable found, just add to the bottom content += ( "\n" + self.WRITE_PATTERN.format(self._name, value) + "\n" ) else: content = self._pattern.sub( self.WRITE_PATTERN.format(self._name, value), content ) atomic_rewrite( self._filename, content, allow_empty_content=self.ALLOW_EMPTY_CONFIG, ) return old_value def get(self) -> Optional[str]: with open(self._filename) as f: content = f.read() return self._parse(content) def _parse(self, content) -> Optional[str]: match = self._pattern.search(content) return match and match.group(1)