majuna/app/terraform/proxy/__init__.py

from abc import abstractmethod
from collections import defaultdict
import datetime
import math
import string
import random
from typing import Dict, Optional, Any, List

from sqlalchemy import text
from tldextract import tldextract

from app import app
from app.extensions import db
from app.models.base import Group
from app.models.mirrors import Proxy
from app.terraform.terraform import TerraformAutomation


class ProxyAutomation(TerraformAutomation):
    subgroup_max = math.inf
    """
    Maximum number of proxies to deploy per sub-group. This is required for some providers
    where the number origins per group may exceed the number of proxies that can be configured
    in a single "configuration block", e.g. Azure CDN's profiles.
    """

    template: str
    """
    Terraform configuration template using Jinja 2.
    """

    template_parameters: List[str]
    """
    List of parameters to be read from the application configuration for use
    in the templating of the Terraform configuration.
    """

    def get_subgroups(self) -> Dict[int, Dict[int, int]]:
        conn = db.engine.connect()
        result = conn.execute(text("""
        SELECT origin.group_id, proxy.psg, COUNT(proxy.id) FROM proxy, origin
        WHERE proxy.origin_id = origin.id
            AND proxy.destroyed IS NULL
            AND proxy.provider = :provider
        GROUP BY origin.group_id, proxy.psg;
        """), provider=self.provider)
        subgroups: Dict[int, Dict[int, int]] = defaultdict(lambda: defaultdict(lambda: 0))
        for row in result:
            subgroups[row[0]][row[1]] = row[2]
        return subgroups

    def create_missing_proxies(self) -> None:
        groups = Group.query.all()
        subgroups = self.get_subgroups()
        for group in groups:
            subgroup = 0
            for origin in group.origins:
                if origin.destroyed is not None:
                    continue
                while True:
                    if subgroups[group.id][subgroup] >= self.subgroup_max:
                        subgroup += 1
                    else:
                        break
                proxies = [
                    x for x in origin.proxies
                    if x.provider == self.provider and x.deprecated is None and x.destroyed is None
                ]
                if not proxies:
                    subgroups[group.id][subgroup] += 1
                    proxy = Proxy()
                    proxy.origin_id = origin.id
                    proxy.provider = self.provider
                    proxy.psg = subgroup
                    # The random usage below is good enough for its purpose: to create a slug that
                    # hasn't been used before.
                    proxy.slug = tldextract.extract(origin.domain_name).domain[:5] + ''.join(
                        random.choices(string.ascii_lowercase, k=12))  # nosec
                    proxy.added = datetime.datetime.utcnow()
                    proxy.updated = datetime.datetime.utcnow()
                    db.session.add(proxy)
        db.session.commit()

    def deprecate_orphaned_proxies(self) -> None:
        proxies = Proxy.query.filter(
            Proxy.deprecated.is_(None),
            Proxy.destroyed.is_(None),
            Proxy.provider == self.provider
        ).all()
        for proxy in proxies:
            if proxy.origin.destroyed is not None:
                proxy.deprecate(reason="origin_destroyed")
        db.session.commit()

    def destroy_expired_proxies(self) -> None:
        cutoff = datetime.datetime.utcnow() - datetime.timedelta(days=3)
        proxies = Proxy.query.filter(
            Proxy.destroyed.is_(None),
            Proxy.provider == self.provider,
            Proxy.deprecated < cutoff
        ).all()
        for proxy in proxies:
            proxy.destroyed = datetime.datetime.utcnow()
            proxy.updated = datetime.datetime.utcnow()
        db.session.commit()

    @abstractmethod
    def import_state(self, state: Any) -> None:
        raise NotImplementedError()

    def tf_prehook(self) -> Optional[Any]:
        self.create_missing_proxies()
        self.deprecate_orphaned_proxies()
        self.destroy_expired_proxies()
        return None

    def tf_posthook(self, *, prehook_result: Any = None) -> None:
        self.import_state(self.tf_show())

    def tf_generate(self) -> None:
        self.tf_write(
            self.template,
            groups=Group.query.all(),
            proxies=Proxy.query.filter(
                Proxy.provider == self.provider,
                Proxy.destroyed.is_(None)
            ).all(),
            subgroups=self.get_subgroups(),
            global_namespace=app.config['GLOBAL_NAMESPACE'],
            bypass_token=app.config['BYPASS_TOKEN'],
            **{
                k: app.config[k.upper()]
                for k in self.template_parameters
            }
        )