majuna/app/cli/automate.py

import logging
import os
import shutil
import tempfile
from datetime import datetime, timedelta, timezone
from traceback import TracebackException
from typing import Type

from app import app
from app.cli import BaseCliHandler, _SubparserType
from app.extensions import db
from app.models.activity import Activity
from app.models.automation import Automation, AutomationLogs, AutomationState
from app.terraform import BaseAutomation
from app.terraform.alarms.eotk_aws import AlarmEotkAwsAutomation
from app.terraform.alarms.proxy_azure_cdn import AlarmProxyAzureCdnAutomation
from app.terraform.alarms.proxy_cloudfront import AlarmProxyCloudfrontAutomation
from app.terraform.alarms.proxy_http_status import AlarmProxyHTTPStatusAutomation
from app.terraform.alarms.smart_aws import AlarmSmartAwsAutomation
from app.terraform.block.block_blocky import BlockBlockyAutomation
from app.terraform.block.block_scriptzteam import BlockBridgeScriptzteamAutomation
from app.terraform.block.bridge_github import BlockBridgeGitHubAutomation
from app.terraform.block.bridge_gitlab import BlockBridgeGitlabAutomation
from app.terraform.block.bridge_roskomsvoboda import BlockBridgeRoskomsvobodaAutomation
from app.terraform.block_external import BlockExternalAutomation
from app.terraform.block_ooni import BlockOONIAutomation
from app.terraform.block_roskomsvoboda import BlockRoskomsvobodaAutomation
from app.terraform.bridge.aws import BridgeAWSAutomation
from app.terraform.bridge.gandi import BridgeGandiAutomation
from app.terraform.bridge.hcloud import BridgeHcloudAutomation
from app.terraform.bridge.meta import BridgeMetaAutomation
from app.terraform.bridge.ovh import BridgeOvhAutomation
from app.terraform.eotk.aws import EotkAWSAutomation
from app.terraform.list.github import ListGithubAutomation
from app.terraform.list.gitlab import ListGitlabAutomation
from app.terraform.list.http_post import ListHttpPostAutomation
from app.terraform.list.s3 import ListS3Automation
from app.terraform.proxy.azure_cdn import ProxyAzureCdnAutomation
from app.terraform.proxy.cloudfront import ProxyCloudfrontAutomation
from app.terraform.proxy.fastly import ProxyFastlyAutomation
from app.terraform.proxy.meta import ProxyMetaAutomation
from app.terraform.static.aws import StaticAWSAutomation
from app.terraform.static.meta import StaticMetaAutomation

jobs = {
    x.short_name: x  # type: ignore[attr-defined]
    for x in [
        # Check for blocked resources first
        BlockBridgeGitHubAutomation,
        BlockBridgeGitlabAutomation,
        BlockBridgeRoskomsvobodaAutomation,
        BlockBridgeScriptzteamAutomation,
        BlockBlockyAutomation,
        BlockExternalAutomation,
        BlockOONIAutomation,
        BlockRoskomsvobodaAutomation,
        # Create new resources
        BridgeMetaAutomation,
        StaticMetaAutomation,
        ProxyMetaAutomation,
        # Terraform
        BridgeAWSAutomation,
        BridgeGandiAutomation,
        BridgeHcloudAutomation,
        BridgeOvhAutomation,
        StaticAWSAutomation,
        EotkAWSAutomation,
        ProxyAzureCdnAutomation,
        ProxyCloudfrontAutomation,
        ProxyFastlyAutomation,
        # Import alarms
        AlarmEotkAwsAutomation,
        AlarmProxyAzureCdnAutomation,
        AlarmProxyCloudfrontAutomation,
        AlarmProxyHTTPStatusAutomation,
        AlarmSmartAwsAutomation,
        # Update lists
        ListGithubAutomation,
        ListGitlabAutomation,
        ListHttpPostAutomation,
        ListS3Automation,
    ]
}


def run_all(**kwargs: bool) -> None:
    """
    Run all automation tasks.

    :param kwargs: this function takes the same arguments as :func:`run_job` and will pass the same arguments
                   to every task
    :return: None
    """
    for job in jobs.values():
        run_job(job, **kwargs)


def run_job(
    job_cls: Type[BaseAutomation], *, force: bool = False, ignore_schedule: bool = False
) -> None:
    automation = Automation.query.filter(
        Automation.short_name == job_cls.short_name
    ).first()
    if automation is None:
        automation = Automation()
        automation.short_name = job_cls.short_name
        automation.description = job_cls.description
        automation.enabled = False
        automation.next_is_full = False
        automation.added = datetime.now(tz=timezone.utc)
        automation.updated = automation.added
        db.session.add(automation)
        db.session.commit()
    else:
        if automation.state == AutomationState.RUNNING and not force:
            logging.warning("Not running an already running automation")
            return
        if not ignore_schedule and not force:
            if automation.next_run is not None and automation.next_run > datetime.now(
                tz=timezone.utc
            ):
                logging.warning("Not time to run this job yet")
                return
    if not automation.enabled and not force:
        logging.warning(
            "job %s is disabled and --force not specified", job_cls.short_name
        )
        return
    automation.state = AutomationState.RUNNING
    db.session.commit()
    try:
        if "TERRAFORM_DIRECTORY" in app.config:
            working_dir = os.path.join(
                app.config["TERRAFORM_DIRECTORY"],
                job_cls.short_name or job_cls.__class__.__name__.lower(),
            )
        else:
            working_dir = tempfile.mkdtemp()
        job: BaseAutomation = job_cls(working_dir)
        success, logs = job.automate()
        # We want to catch any and all exceptions that would cause problems here, because
        # the error handling process isn't really handling the error, but rather causing it
        # to be logged for investigation. Catching more specific exceptions would just mean that
        # others go unrecorded and are difficult to debug.
    except Exception as exc:  # pylint: disable=broad-except
        if logging.getLogger().level == logging.DEBUG:
            raise exc
        trace = TracebackException.from_exception(exc)
        success = False
        logs = "\n".join(trace.format())
    if job is not None and success:
        automation.state = AutomationState.IDLE
        automation.next_run = datetime.now(tz=timezone.utc) + timedelta(
            minutes=getattr(job, "frequency", 7)
        )
        if "TERRAFORM_DIRECTORY" not in app.config and working_dir is not None:
            # We used a temporary working directory
            shutil.rmtree(working_dir)
    else:
        automation.state = AutomationState.ERROR
        safe_jobs = [
            "proxy_cloudfront",
            "static_aws",
            "list_http_post",
            "list_s3",
            "list_github",
            "list_gitlab",
            "block_blocky",
            "block_external",
            "block_ooni",
        ]
        if job.short_name not in safe_jobs:
            automation.enabled = False
        automation.next_run = None
        log = AutomationLogs()
        log.automation_id = automation.id
        log.added = datetime.now(tz=timezone.utc)
        log.updated = datetime.now(tz=timezone.utc)
        log.logs = str(logs)
        db.session.add(log)
        db.session.commit()
        activity = Activity(
            activity_type="automation",
            text=(
                f"[{automation.short_name}] 🚨 Automation failure: It was not possible to handle this failure safely "
                "and so the automation task has been automatically disabled. It may be possible to simply re-enable "
                "the task, but repeated failures will usually require deeper investigation. See logs for full "
                "details."
            ),
        )
        db.session.add(activity)
        activity.notify()  # Notify before commit because the failure occurred even if we can't commit.
    automation.last_run = datetime.now(tz=timezone.utc)
    db.session.commit()


class AutomateCliHandler(BaseCliHandler):
    @classmethod
    def add_subparser_to(cls, subparsers: _SubparserType) -> None:
        parser = subparsers.add_parser("automate", help="automation operations")
        parser.add_argument(
            "-a",
            "--all",
            dest="all",
            help="run all automation jobs",
            action="store_true",
        )
        parser.add_argument(
            "-j",
            "--job",
            dest="job",
            choices=sorted(jobs.keys()),
            help="run a specific automation job",
        )
        parser.add_argument(
            "--force",
            help="run job even if disabled and it's not time yet",
            action="store_true",
        )
        parser.add_argument(
            "--ignore-schedule",
            help="run job even if it's not time yet",
            action="store_true",
        )
        parser.set_defaults(cls=cls)

    def run(self) -> None:
        with app.app_context():
            if self.args.job:
                run_job(
                    jobs[self.args.job],
                    force=self.args.force,
                    ignore_schedule=self.args.ignore_schedule,
                )
            elif self.args.all:
                run_all(
                    force=self.args.force, ignore_schedule=self.args.ignore_schedule
                )
            else:
                logging.error("No action requested")