majuna/app/terraform/proxy/__init__.py

from collections import defaultdict
import datetime
import math
import string
import random

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

    def get_subgroups(self):
        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 = defaultdict(lambda: defaultdict(lambda: 0))
        for row in result:
            subgroups[row[0]][row[1]] = row[2]
        return subgroups

    def create_missing_proxies(self):
        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
                    proxy.slug = tldextract.extract(origin.domain_name).domain[:5] + ''.join(
                        random.choices(string.ascii_lowercase, k=12))
                    proxy.added = datetime.datetime.utcnow()
                    proxy.updated = datetime.datetime.utcnow()
                    db.session.add(proxy)
        db.session.commit()

    def deprecate_orphaned_proxies(self):
        proxies = Proxy.query.filter(
            Proxy.deprecated == None,
            Proxy.destroyed == 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):
        cutoff = datetime.datetime.utcnow() - datetime.timedelta(days=3)
        proxies = Proxy.query.filter(
            Proxy.destroyed == 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()

    def tf_prehook(self):
        self.create_missing_proxies()
        self.deprecate_orphaned_proxies()
        self.destroy_expired_proxies()

    def tf_posthook(self, *, prehook_result):
        self.import_state(self.tf_show())

    def tf_generate(self):
        self.tf_write(
            self.template,
            groups=Group.query.all(),
            proxies=Proxy.query.filter(
                Proxy.provider == self.provider,
                Proxy.destroyed == 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
            }
        )