def update_traffic(
server: Server, traffic: str, accumulate: bool = False
):
pattern = re.compile(r"\/\* (UPLOAD|DOWNLOAD)(?:\-UDP)? ([0-9]+)->")
prev_ports = {port.num: port for port in server.ports}
db_ports = {}
traffics = defaultdict(lambda: {"download": 0, "upload": 0})
for line in traffic.split("\n"):
match = pattern.search(line)
if match and len(match.groups()) > 1 and match.groups()[1].isdigit():
port_num = int(match.groups()[1])
traffics[port_num][match.groups()[0].lower()] += int(
line.split()[1]
)
with db_session() as db:
for port_num, usage in traffics.items():
update_usage(
db, prev_ports, db_ports, server.id, port_num, usage, accumulate
)
server_users_usage = defaultdict(lambda: {"download": 0, "upload": 0})
with db_session() as db:
server = get_server_with_ports_usage(db, server.id)
for port in server.ports:
if port.usage:
check_port_limits(db, port)
for port_user in port.allowed_users:
server_users_usage[port_user.user_id][
"download"
] += port.usage.download
server_users_usage[port_user.user_id][
"upload"
] += port.usage.upload
check_server_user_limit(db, server, server_users_usage)
def iptables_runner(
port_id: int,
server_id: int,
local_port: int,
remote_address: str,
remote_port: int = None,
forward_type: str = None,
update_status: bool = False,
):
try:
if not is_ip(remote_address):
remote_ip = dns_query(remote_address)
else:
remote_ip = remote_address
with db_session() as db:
port = get_port(db, server_id, port_id)
if not forward_type:
args = f" delete {local_port}"
elif remote_port:
port.forward_rule.config["remote_ip"] = remote_ip
db.add(port.forward_rule)
db.commit()
args = (
f" -t={forward_type}"
f" {'-v=6' if is_ipv6(remote_ip) else '-v=4'}"
f" forward {local_port} {remote_ip} {remote_port}"
)
else:
args = f" list {local_port}"
server = get_server_with_ports_usage(db, server_id)
extravars = {
"host": server.ansible_name,
"local_port": local_port,
"iptables_args": args,
}
run(
server=server,
playbook="iptables.yml",
extravars=extravars,
status_handler=lambda s, **k: status_handler(
port_id, s, update_status
),
finished_callback=iptables_finished_handler(
server.id, port_id, True
)
if update_status
else lambda r: None,
)
except Exception:
traceback.print_exc()
with db_session() as db:
port = get_port(db, server_id, port_id)
port.forward_rule.status = "failed"
port.forward_rule.config["error"] = traceback.format_exc()
print(port.forward_rule.__dict__)
db.add(port.forward_rule)
db.commit()
def clean_port_no_update_runner(server_id: int, port_num: int):
with db_session() as db:
server = get_server_with_ports_usage(db, server_id)
run(
server=server,
playbook="clean_port.yml",
extravars={"local_port": port_num},
)
def traffic_server_runner(server_id: Server):
with db_session() as db:
server = get_server_with_ports_usage(db, server_id)
run(
server=server,
playbook="traffic.yml",
finished_callback=iptables_finished_handler(server.id),
)
def clean_port_runner(server_id: int, port: Port, update_traffic: bool = True):
with db_session() as db:
if db_forward_rule := get_forward_rule(db, server_id, port.id):
db.delete(db_forward_rule)
db.commit()
server = get_server_with_ports_usage(db, server_id)
if db_forward_rule := get_forward_rule(db, server_id, port.id):
db.delete(db_forward_rule)
db.commit()
def app_runner(
port_id: int,
server_id: int,
port_num: int,
app_name: str,
app_command: str = None,
app_config: t.Dict = None,
app_version_arg: str = "-v",
traffic_meter: bool = True,
app_role_name: str = "app",
app_download_role_name: str = None,
app_sync_role_name: str = "app_sync",
app_get_role_name: str = "app_get",
remote_ip: str = "ANYWHERE",
ident: str = None,
update_status: bool = False,
):
with db_session() as db:
server = get_server_with_ports_usage(db, server_id)
extravars = {
"host": server.ansible_name,
"local_port": port_num,
"remote_ip": remote_ip,
"app_name": app_name,
"app_command": app_command,
"app_version_arg": app_version_arg,
"traffic_meter": traffic_meter,
"app_download_role_name": app_download_role_name
if app_download_role_name is not None
else f"{app_name}_download",
"app_role_name": app_role_name,
"app_sync_role_name": app_sync_role_name,
"app_get_role_name": app_get_role_name,
"update_status": update_status,
"update_app": update_status and not server.config.get(app_name),
}
if app_config is not None:
with open(
f"ansible/project/roles/app/files/{app_name}-{port_id}", "w"
) as f:
f.write(app_config)
extravars["app_config"] = f"{app_name}-{port_id}"
run(
server=server,
playbook="app.yml",
extravars=extravars,
ident=ident,
status_handler=lambda s, **k: status_handler(port_id, s, update_status),
finished_callback=iptables_finished_handler(server.id, port_id, True)
if update_status
else lambda r: None,
)
def clean_port_runner(
server_id: int, port_num: int, update_traffic: bool = True
):
with db_session() as db:
server = get_server_with_ports_usage(db, server_id)
run(
server=server,
playbook="clean_port.yml",
extravars={"local_port": port_num},
finished_callback=iptables_finished_handler(
server, accumulate=True, update_traffic_bool=update_traffic
),
)
def rule_runner(rule_id: int):
try:
with db_session() as db:
rule = get_forward_rule_by_id(db, rule_id)
port, port_id, port_num, server_id = (
rule.port,
rule.port.id,
rule.port.num,
rule.port.server.id,
)
ident = uuid4()
app_configs = []
if rule.config.get("reverse_proxy"):
reverse_proxy_port = get_port_by_id(
db, rule.config.get("reverse_proxy")
)
app_configs.append(
AppConfig.configs[
reverse_proxy_port.forward_rule.method
].apply(db, reverse_proxy_port)
)
app_configs.append(
AppConfig.configs[rule.method].apply(db, rule.port)
)
db.refresh(rule)
server = get_server_with_ports_usage(db, server_id)
for config in app_configs:
runner = run(
server,
config.playbook,
extravars=config.extravars,
ident=ident,
status_handler=lambda s, **k: status_handler(port_id, s, True),
finished_callback=iptables_finished_handler(
server.id, port_id, True
),
)
if runner.status != "successful":
break
except Exception:
with db_session() as db:
rule.status = "failed"
rule.config["error"] = traceback.format_exc()
db.add(rule)
db.commit()
async def root():
with db_session() as db:
server = get_server_with_ports_usage(db, 34)
print([p for p in server.ports])
return {"message": "Hello World"}
