majuna/app/cli/automate.py

import datetime
import logging
import os
import shutil
import tempfile
from traceback import TracebackException
from typing import Type

from app import app
from app.cli import _SubparserType, BaseCliHandler
from app.extensions import db
from app.models.activity import Activity
from app.models.automation import Automation, AutomationState, AutomationLogs
from app.terraform import BaseAutomation
from app.terraform.block.bridge_dnsc import BlockBridgeDnscAutomation
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.meta import BridgeMetaAutomation
from app.terraform.eotk.aws import EotkAWSAutomation
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.bridge.aws import BridgeAWSAutomation
from app.terraform.bridge.gandi import BridgeGandiAutomation
from app.terraform.bridge.hcloud import BridgeHcloudAutomation
from app.terraform.bridge.ovh import BridgeOvhAutomation
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.meta import ProxyMetaAutomation
from app.terraform.proxy.azure_cdn import ProxyAzureCdnAutomation
from app.terraform.proxy.cloudfront import ProxyCloudfrontAutomation
from app.terraform.proxy.fastly import ProxyFastlyAutomation

jobs = {
    x.short_name: x  # type: ignore[attr-defined]
    for x in [
        AlarmEotkAwsAutomation,
        AlarmProxyAzureCdnAutomation,
        AlarmProxyCloudfrontAutomation,
        AlarmProxyHTTPStatusAutomation,
        AlarmSmartAwsAutomation,
        BlockBridgeDnscAutomation,
        BlockBridgeGitHubAutomation,
        BlockBridgeGitlabAutomation,
        BlockBridgeRoskomsvobodaAutomation,
        BlockExternalAutomation,
        BlockOONIAutomation,
        BlockRoskomsvobodaAutomation,
        BridgeAWSAutomation,
        BridgeGandiAutomation,
        BridgeHcloudAutomation,
        BridgeMetaAutomation,
        BridgeOvhAutomation,
        EotkAWSAutomation,
        ListGithubAutomation,
        ListGitlabAutomation,
        ListHttpPostAutomation,
        ListS3Automation,
        ProxyAzureCdnAutomation,
        ProxyCloudfrontAutomation,
        ProxyFastlyAutomation,
        ProxyMetaAutomation
    ]
}


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.datetime.utcnow()
        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.datetime.utcnow():
                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
        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.datetime.utcnow() + datetime.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
        automation.enabled = False
        automation.next_run = None
        log = AutomationLogs()
        log.automation_id = automation.id
        log.added = datetime.datetime.utcnow()
        log.updated = datetime.datetime.utcnow()
        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.datetime.utcnow()
    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")