Implement disabled checks

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)
 

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())

pylibagent/version.py

@@ -1 +1 @@
-__version__ = '0.3.1'
+__version__ = '1.0.0'