111 lines
3.8 KiB
Python
111 lines
3.8 KiB
Python
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
|
|
}
|
|
)
|