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 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.
    """

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

    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 }}.{{ provider }}.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("provider", "equalto", provider) | 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()
        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=cls.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

    @classmethod
    def next_subgroup(cls, group_id: int) -> Optional[int]:
        conn = db.engine.connect()
        result = conn.execute(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;
        """), provider=cls.short_name, group_id=group_id)
        subgroups = {
            row[0]: row[1] for row in result
        }
        for subgroup in range(0, cls.subgroup_count_max):
            if subgroups.get(subgroup, 0) < cls.subgroup_members_max:
                return subgroup
        return None