1#@!#!123s

: / opt / imunify360 / venv / lib / python3.11 / site-packages / im360 / plugins / resident /

Filename : health.py

back

import json
import time

from defence360agent.api import health
from defence360agent.contracts.messages import MessageType
from defence360agent.contracts.plugins import (
    MessageSink,
    MessageSource,
    expect,
)
from defence360agent.rpc_tools import ValidationError


class HealthChecker(MessageSink):
    async def create_sink(self, loop):
        self._loop = loop

    @expect(MessageType.Health)
    async def check_health(self, message):
        transport = message.get("transport", None)
        try:
            response = health.sensor.status(time.time())._asdict()
        except Exception as ex:
            self._respond_to_transport(
                transport, json.dumps({"error": str(ex)})
            )
        self._respond_to_transport(transport, json.dumps(response))

    def _respond_to_transport(self, transport, message):
        """Writes the provided message to the transport if it exists"""
        if transport:
            transport.write(message.encode() + b"\n")