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()