majuna/app/terraform/proxy/__init__.py

import os.path
import sys
from abc import abstractmethod
import datetime
from collections import defaultdict
from typing import Optional, Any, List, Dict

from flask import current_app
from sqlalchemy import text

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


def update_smart_proxy_instance(group_id: int,
                                provider: str,
                                region: str,
                                instance_id: str) -> None:
    instance = SmartProxy.query.filter(
        SmartProxy.group_id == group_id,
        SmartProxy.region == region,
        SmartProxy.provider == provider,
        SmartProxy.destroyed.is_(None)
    ).first()
    if instance is None:
        instance = SmartProxy()
        instance.added = datetime.datetime.utcnow()
        instance.group_id = group_id
        instance.provider = provider
        instance.region = region
        db.session.add(instance)
    instance.updated = datetime.datetime.utcnow()
    instance.instance_id = instance_id


class ProxyAutomation(TerraformAutomation):
    subgroup_members_max = sys.maxsize
    """
    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.
    """

    subgroup_count_max = sys.maxsize
    """
    Maximum number of subgroups that can be deployed. This is required for some providers where
    the total number of subgroups is limited by a quota, 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.
    """

    smart_proxies = False
    """
    Whether this provider supports "smart" proxies.
    """

    cloud_name: str
    """
    The name of the cloud provider used by this proxy provider. Used to determine if this provider can be enabled.
    """

    provider: str  # type: ignore[assignment]
    # TODO: Temporary override

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

    def enabled(self) -> bool:
        return current_app.config.get(self.cloud_name.upper() + "_ENABLED", False)

    def tf_prehook(self) -> Optional[Any]:  # pylint: disable=useless-return
        return None

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

    def tf_generate(self) -> None:
        groups = Group.query.all()
        self.tf_write(
            self.template,
            groups=groups,
            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'],
            terraform_modules_path=os.path.join(*list(os.path.split(app.root_path))[:-1], 'terraform-modules'),
            backend_config=f"""backend "http" {{
              lock_address = "{app.config['TFSTATE_BACKEND']}/{self.short_name}"
              unlock_address = "{app.config['TFSTATE_BACKEND']}/{self.short_name}"
              address = "{app.config['TFSTATE_BACKEND']}/{self.short_name}"
            }}""",
            **{k: app.config[k.upper()] for k in self.template_parameters})
        if self.smart_proxies:
            for group in groups:
                self.sp_config(group)

    def sp_config(self, group: Group) -> None:
        group_origins: List[Origin] = Origin.query.filter(
            Origin.group_id == group.id,
            Origin.destroyed.is_(None),
            Origin.smart.is_(True)
        ).all()
        self.tmpl_write(f"smart_proxy.{group.id}.conf", """
        {% for origin in origins %}
        server {
          listen 443 ssl;
          server_name origin-{{ origin.id }}.{{ origin.group.group_name | lower }}.smart.{{ smart_zone[:-1] }};
          location / {
            proxy_set_header Accept-Encoding "";
            proxy_ssl_server_name on;
            proxy_pass https://{{ origin.domain_name }}/;
            subs_filter_types text/html text/css text/xml;
            subs_filter https://{{ origin.domain_name }}/ /;
            subs_filter "([^:]|)\\\"https://{{ origin.domain_name }}\\\"" \\1\\\"/\\\";
            {%- for asset_origin in origin.group.origins | selectattr("assets") -%}
            {%- for asset_proxy in asset_origin.proxies | selectattr("deprecated", "none") | selectattr("destroyed", "none") -%}
            {%- if loop.first %}
            subs_filter https://{{ asset_origin.domain_name }}/ {{ asset_proxy.url }}/;
            {%- endif -%}
            {%- endfor -%}
            {%- endfor %}
          }
          ssl_certificate /etc/ssl/smart_proxy.crt;
          ssl_certificate_key /etc/ssl/private/smart_proxy.key;
        }
        {% endfor %}
        """,
                        provider=self.provider,
                        origins=group_origins,
                        smart_zone=app.config['SMART_ZONE'])

    @classmethod
    def get_subgroups(cls) -> Dict[int, Dict[int, int]]:
        conn = db.engine.connect()
        stmt = 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;
        """)
        stmt = stmt.bindparams(provider=cls.provider)
        result = conn.execute(stmt).all()
        subgroups: Dict[int, Dict[int, int]] = defaultdict(lambda: defaultdict(lambda: 0))
        for row in result:
            subgroups[row[0]][row[1]] = row[2]
        return subgroups

    @classmethod
    def next_subgroup(cls, group_id: int) -> Optional[int]:
        provider = cls()  # Some attributes are set by constructor
        conn = db.engine.connect()
        stmt = text("""
        SELECT proxy.psg, COUNT(proxy.id) FROM proxy, origin
        WHERE proxy.origin_id = origin.id
            AND proxy.destroyed IS NULL
            AND origin.group_id = :group_id
            AND proxy.provider = :provider
        GROUP BY proxy.psg ORDER BY proxy.psg;
        """)
        stmt = stmt.bindparams(provider=provider.short_name, group_id=group_id)
        result = conn.execute(stmt).all()
        subgroups = {
            row[0]: row[1] for row in result
        }
        for subgroup in range(0, provider.subgroup_count_max):
            if subgroups.get(subgroup, 0) < provider.subgroup_members_max:
                return subgroup
        return None