utils.py

"""Utilities for common checks."""

import os
import boto3  # pylint: disable=import-error
import requests  # pylint: disable=import-error,useless-suppression


class Check:
    _topic_arn = os.getenv("TOPIC_ARN")
    _retries = int(os.getenv("RETRIES", "2"))

    def run(self):
        """Run the check function with retries. Alert if failed."""
        for _ in range(self._retries):
            success, message = self._check()
            print(message)
            if success:
                break
        else:
            self.publish(message)
        return success

    def _check(self):
        """The actual check."""
        raise NotImplementedError

    def publish(self, message):
        """Publish an SNS message."""
        if self._topic_arn is None:
            print(f"Publish: {message}")
        else:
            client = boto3.client("sns")
            client.publish(TopicArn=self._topic_arn, Message=message)


class CheckURL(Check):
    _method = "GET"
    _valid_codes = (200,)

    def __init__(self, url, method=None, codes=None):
        """Initialize the check object."""
        self._url = url
        if method is not None:
            self._method = method
        if codes is not None:
            self._valid_codes = codes

    def _check(self):
        """Checks validaty of a URL."""
        try:
            response = requests.request(
                self._method, self._url, allow_redirects=False
            )
            return (
                response.status_code in self._valid_codes,
                f"{self._url} is OK.",
            )
        except Exception as e:  # pylint: disable=broad-except,invalid-name
            return False, str(e)


def website_handler(urls):
    """Return a Lambda event handler for checking web sites."""

    def _handler(event, context):  # pylint: disable=unused-argument
        for url in urls:
            CheckURL(url["url"], url.get("method"), url.get("codes")).run()

    return _handler