135 lines
4.8 KiB
Python
135 lines
4.8 KiB
Python
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 == 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) -> None:
|
|
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()
|
|
|
|
@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 == 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
|
|
}
|
|
)
|