from collections import defaultdict import datetime import math import string import random from typing import Dict 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) -> 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 = 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 } )