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 ddns_runner():
with db_session() as db:
rules = get_all_ddns_rules(db)
for rule in rules:
remote_addr = ''
if rule.config.get("remote_address"):
remote_addr = rule.config.get("remote_address").strip(' ')
if (
rule.config.get("remote_address")
and rule.config.get("remote_ip")
and not is_ip(remote_addr)
):
updated_ip = dns_query(remote_addr)
if updated_ip and updated_ip != rule.config["remote_ip"]:
print(
f"DNS changed for address {remote_addr}, "
+ f"{rule.config['remote_ip']}->{updated_ip}"
)
if rule.method == MethodEnum.IPTABLES:
iptables_runner(
rule.port.id,
rule.port.server.id,
rule.port.num,
remote_address=updated_ip,
remote_port=rule.config["remote_port"],
forward_type=rule.config.get("type", "ALL"),
update_status=True,
)
else:
rule_runner(rule_id=rule.id)
def servers_runner(**kwargs):
with db_session() as db:
servers = get_servers(db)
init_md5 = get_md5_for_file("ansible/project/server.yml")
for server in servers:
if "init" not in server.config or server.config["init"] != init_md5:
server_runner.delay(server.id, **kwargs)
def connect_runner(server_id: int, ):
with db_session() as db:
server = get_server(db, server_id)
return run(
server=server,
playbook="connect.yml",
finished_callback=finished_handler(server),
)
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 db_session_middleware(request: Request, call_next):
"""
中间件，在每次收到的 request 时，生成一个数据库会话并保存在 request.state 中，并确保在 response 前关闭
"""
request.state.db = db_session()
response = await call_next(request)
request.state.db.close()
return response
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_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},
)
async def db_session_middleware(request: Request, call_next):
with db_session() as db:
request.state.db = db
try:
response = await call_next(request)
return response
except IntegrityError as e:
return JSONResponse(status_code=400,
content={"detail": str(e.orig)})
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 connect_runner(server_id: int, ):
with db_session() as db:
server = get_server(db, server_id)
run(
server=server,
playbook="connect.yml",
event_handler=server_facts_event_handler(server.id),
finished_callback=finished_handler(server.id),
)
def clean_expired_port_runner():
with db_session() as db:
db_expire_rules = get_all_expire_rules(db)
for db_rule in db_expire_rules:
if time.time() > db_rule.config.get("expire_time", float("inf")):
clean_port_runner(
db_rule.port.server.id,
db_rule.port,
update_traffic=True,
)
def update_rule_error(server_id: int, port_id: int, facts: t.Dict):
with db_session() as db:
db_rule = get_forward_rule(db, server_id, port_id)
db_rule.config["error"] = "\n".join([facts.get("error", "")] + [
re.search(r"\w+\[[0-9]+\]: (.*)$", line).group(1)
for line in facts.get("systemd_error", "").split("\n")
if re.search(r"\w+\[[0-9]+\]: (.*)$", line)
]).strip()
db.add(db_rule)
db.commit()
def server_runner(server_id: int, **kwargs):
init_md5 = get_md5_for_file("ansible/project/server.yml")
with db_session() as db:
server = get_server(db, server_id)
run(
server=server,
playbook="server.yml",
extravars=kwargs,
event_handler=server_facts_event_handler(server.id),
finished_callback=finished_handler(server, init_md5),
)
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 ansible_hosts_runner():
with db_session() as db:
servers = db.query(Server).filter(Server.is_active == True).all()
with open("ansible/inventory/hosts", 'w+') as f:
f.write("### START AUTO GENERATION ###\n")
for server in servers:
f.write(f"{server.ansible_name}"
f"\tansible_host={server.ansible_host}"
f"\tansible_port={server.ansible_port}"
f"\tansible_user={server.ansible_user}\n")
f.write("### END AUTO GENERATION ###")
def wrapper(runner):
with db_session() as db:
server = get_server(db, server_id)
facts = runner.get_fact_cache(server.ansible_name)
if facts:
if facts.get("traffic", "") and update_traffic_bool:
update_traffic(server,
facts.get("traffic", ""),
accumulate=accumulate)
if port_id is not None and (facts.get("error")
or facts.get("systemd_error")):
update_rule_error(server.id, port_id, facts)
update_facts(server.id, facts)
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 status_handler(port_id: int, status_data: dict, update_status: bool):
if update_status:
with db_session() as db:
rule = (db.query(PortForwardRule).filter(
PortForwardRule.port_id == port_id).first())
if rule:
if (status_data.get("status", None) == "starting"
and rule.status == "running"):
return status_data
if status_data.get("runner_ident"):
rule.config["runner"] = status_data.get("runner_ident")
rule.status = status_data.get("status", None)
db.add(rule)
db.commit()
return status_data
def iptables_reset_runner(
server_id: int,
port_num: int,
):
with db_session() as db:
server = get_server(db, server_id)
extravars = {
"host": server.ansible_name,
"local_port": port_num,
"iptables_args": f" reset {port_num}",
}
run(
server=server,
playbook="iptables.yml",
extravars=extravars,
)
def tc_runner(
server_id: int,
port_num: int,
egress_limit: int = None,
ingress_limit: int = None
):
with db_session() as db:
server = get_server(db, server_id)
args = ""
if egress_limit:
args += f' -e={egress_limit}kbit'
if ingress_limit:
args += f' -i={ingress_limit}kbit'
args += f' {port_num}'
run(
server=server,
playbook="tc.yml",
extravars={"host": server.ansible_name, "tc_args": args},
)
def crawl_and_create(db_data):
if db_data.desc != 'null':
return False
if db_data.play_url == '':
return False
site = db_data.sid
while True:
detail = crawler_list[site - 1].crawl_detail(db_data)
if detail:
break
db = db_session()
crud_crawler.create_detail(db=db,
detail={
'aid': detail['aid'],
'desc': detail['desc']
},
director=detail['director'],
actor=detail['actor'])
db.close()
return True
def init() -> None:
email = input("请输入管理员邮箱：")
if not email or "@" not in email:
print(" 邮箱格式不正确 ！")
return
password = getpass("请输入密码：")
repeated_password = getpass("请再次输入密码：")
if password != repeated_password:
print("两次密码不一致！")
return
with db_session() as db:
create_user(
db,
UserCreate(
email=email,
password=password,
is_active=True,
is_superuser=True,
),
)
def update_facts(server_id: int, facts: t.Dict, md5: str = None):
with db_session() as db:
db_server = get_server(db, server_id)
if facts.get("ansible_os_family"):
db_server.config["system"] = {
"os_family": facts.get("ansible_os_family"),
"architecture": facts.get("ansible_architecture"),
"distribution": facts.get("ansible_distribution"),
"distribution_version":
facts.get("ansible_distribution_version"),
"distribution_release":
facts.get("ansible_distribution_release"),
}
elif facts.get("msg"):
db_server.config["system"] = {"msg": facts.get("msg")}
if "services" in facts:
db_server.config["services"] = facts.get("services")
# TODO: Add disable feature
for func in [
"brook",
"caddy",
"ehco",
"gost",
"iperf",
"iptables",
"node_exporter",
"shadowsocks",
"socat",
"tiny_port_mapper",
"v2ray",
"wstunnel",
"realm",
"haproxy",
]:
if func in facts:
db_server.config[func] = facts.get(func)
if md5 is not None:
db_server.config["init"] = md5
db.add(db_server)
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"}
def traffic_runner():
with db_session() as db:
servers = get_servers(db)
for server in servers:
traffic_server_runner.delay(server.id)
def wrapper(runner):
with db_session() as db:
server = get_server(db, server_id)
facts = runner.get_fact_cache(server.ansible_name)
update_facts(server.id, facts, md5=md5)
def traffic_runner():
with db_session() as db:
servers = get_servers(db)
for server in servers:
traffic_server_runner(server.id)
clear_cache()
