diff --git a/README.md b/README.md index dfdb0b3..56cc582 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,7 @@ Variable | Default | Description `LOG_LEVEL` | `warning` | Log level _(error, warning, info, debug)_. `LOG_COLORIZED` | `0` | Log colorized, 0 _(=disabled)_ or 1 _(=enabled)_. `LOG_FMT` | `%y%m...` | Default format is `%y%m%d %H:%M:%S`. - +`DISABLED_CHECKS_CACHE_AGE` | `900` | Time in seconds before we poll again for reading the disabled checks. ## Usage (Demonized agent) diff --git a/pylibagent/agent.py b/pylibagent/agent.py index 08f8b56..ee0104a 100644 --- a/pylibagent/agent.py +++ b/pylibagent/agent.py @@ -65,10 +65,10 @@ def __init__(self, key: str, version: str): logging.error('missing environment variable `TOKEN`') exit(1) + self.headers = {'Authorization': f'Bearer {token}'} self._loop: Optional[asyncio.AbstractEventLoop] = None - self._headers = {'Authorization': f'Bearer {token}'} self._json_headers = {'Content-Type': 'application/json'} - self._json_headers.update(self._headers) + self._json_headers.update(self.headers) self.asset_id: Optional[int] = None self.asset_id_file: str = asset_id_file @@ -79,6 +79,8 @@ def __init__(self, key: str, version: str): else: self._read_json() + self._disabled_checks: DisabledChecks = DisabledChecks() + async def announce(self, asset_name: Optional[str] = None, asset_kind: Optional[str] = None): """Announce the agent. @@ -94,7 +96,7 @@ async def announce(self, asset_name: Optional[str] = None, return url = _join(self.api_uri, f'asset/{self.asset_id}') - async with ClientSession(headers=self._headers) as session: + async with ClientSession(headers=self.headers) as session: async with session.get( url, params={'fields': 'name', 'collectors': 'key'}, @@ -117,7 +119,7 @@ async def announce(self, asset_name: Optional[str] = None, self.api_uri, f'asset/{self.asset_id}/collector/{self.key}') try: - async with ClientSession(headers=self._headers) as session: + async with ClientSession(headers=self.headers) as session: async with session.post( url, ssl=self.verify_ssl @@ -251,9 +253,12 @@ async def _check_loop(self, check: Type[CheckBase]): ts = ts_next try: - check_data = \ - await asyncio.wait_for(check.run(), timeout=timeout) - await self.send_data(check.key, check_data) + if await self._disabled_checks.is_disabled(self, check.key): + logging.debug(f'check {check.key} is disabled') + else: + check_data = \ + await asyncio.wait_for(check.run(), timeout=timeout) + await self.send_data(check.key, check_data) except asyncio.TimeoutError: logging.error(f'check error ({check.key}): timed out') except SendDataException as e: @@ -271,7 +276,7 @@ async def _create_asset(self, asset_name: Optional[str] = None, asset_kind: Optional[str] = None ) -> Tuple[int, str]: url = _join(self.api_uri, 'container/id') - async with ClientSession(headers=self._headers) as session: + async with ClientSession(headers=self.headers) as session: async with session.get( url, ssl=self.verify_ssl @@ -301,7 +306,7 @@ async def _create_asset(self, asset_name: Optional[str] = None, try: url = _join(self.api_uri, f'asset/{asset_id}/collector/{self.key}') - async with ClientSession(headers=self._headers) as session: + async with ClientSession(headers=self.headers) as session: async with session.post( url, ssl=self.verify_ssl @@ -368,3 +373,42 @@ def _dump_json(self): logging.error( f"failed to write file: {self.asset_id_file} ({msg})") exit(1) + + +class DisabledChecks: + def __init__(self): + self._list = [] + self._age: int = 0 + self._max_age: int = int(os.getenv('DISABLED_CHECKS_CACHE_AGE', '900')) + self._lock: asyncio.Lock = asyncio.Lock() + + async def is_disabled(self, agent: Agent, check_key: str) -> bool: + async with self._lock: + if time.time() - self._age > self._max_age: + await self._update_list(agent) + return check_key in self._list + + async def _update_list(self, agent: Agent): + self._list = [] + try: + url = _join(agent.api_uri, f'asset/{agent.asset_id}') + async with ClientSession(headers=agent.headers) as session: + async with session.get( + url, + params={'fields': 'disabledChecks'}, + ssl=agent.verify_ssl + ) as r: + if r.status != 200: + msg = await r.text() + raise Exception(f'{msg} (error code: {r.status})') + resp = await r.json() + disabledChecks = resp["disabledChecks"] + except Exception as e: + logging.error(e) + return + + for dc in disabledChecks: + if dc['collector'] == agent.key: + self._list.append(dc['check']) + + self._age = int(time.time()) diff --git a/pylibagent/version.py b/pylibagent/version.py index e1424ed..1f356cc 100644 --- a/pylibagent/version.py +++ b/pylibagent/version.py @@ -1 +1 @@ -__version__ = '0.3.1' +__version__ = '1.0.0'