4a82ab0b06
Updated the `create_rule` function to handle a list of ports by joining them into a comma-separated string. This enhances flexibility, allowing both single port and multiple port configurations for port forwarding rules.
131 lines
4 KiB
Python
131 lines
4 KiB
Python
import os
|
|
|
|
import aiohttp
|
|
import asyncio
|
|
import socket
|
|
|
|
import configparser
|
|
|
|
cfg = configparser.ConfigParser()
|
|
|
|
import asusrouter.modules.port_forwarding
|
|
from asusrouter import AsusRouter, AsusData
|
|
from api.pelican import PelicanApi
|
|
from flask import Flask, request
|
|
|
|
RULE_PREFIX = "[auto]::"
|
|
|
|
|
|
def create_rule(label, address, ports):
|
|
|
|
if isinstance(ports, list):
|
|
ports = ",".join(ports)
|
|
|
|
return asusrouter.modules.port_forwarding.PortForwardingRule(RULE_PREFIX + label, address, None, "BOTH", None, ports)
|
|
|
|
|
|
def rule_to_string(rule: asusrouter.modules.port_forwarding.PortForwardingRule):
|
|
return f"{rule.name}::{rule.port_external}->{rule.port if rule.port else rule.port_external}@{rule.ip_address}"
|
|
|
|
|
|
def apply_port_changes():
|
|
loop = asyncio.new_event_loop()
|
|
|
|
session = aiohttp.ClientSession(loop=loop)
|
|
|
|
pelican_cfg = cfg["pelican"]
|
|
router_cfg = cfg["router"]
|
|
pelican = PelicanApi(pelican_cfg["baseurl"], pelican_cfg["auth"], pelican_cfg["client_auth"])
|
|
router = AsusRouter(
|
|
hostname=router_cfg["hostname"],
|
|
username=router_cfg["username"],
|
|
password=router_cfg["password"],
|
|
use_ssl=False,
|
|
session=session
|
|
)
|
|
|
|
run = loop.run_until_complete
|
|
|
|
run(router.async_connect())
|
|
|
|
data = run(router.async_get_data(AsusData.PORT_FORWARDING))
|
|
rules = data["rules"]
|
|
x: asusrouter.modules.port_forwarding.PortForwardingRule
|
|
# Remove auto-generated rules
|
|
rules = [x for x in rules if not x.name.startswith(RULE_PREFIX)]
|
|
|
|
# Add new auto-generated rules
|
|
allocation_server_lookup = {}
|
|
servers = pelican.servers()
|
|
for server in servers:
|
|
server_allocations = pelican.server_allocations(server["identifier"])
|
|
|
|
for alloc in server_allocations:
|
|
allocation_server_lookup[alloc["id"]] = {
|
|
"server": server,
|
|
"resources": pelican.resources(server["identifier"])
|
|
}
|
|
|
|
for node in pelican.nodes():
|
|
node_id = node["id"]
|
|
allocs = pelican.allocations(node_id)
|
|
|
|
node_internal_name = node["name"] + ".pve.local"
|
|
node_internal_ip = socket.getaddrinfo(node_internal_name, 0)[0][4][0]
|
|
|
|
for alloc in allocs:
|
|
if not alloc["assigned"]:
|
|
continue
|
|
|
|
label = "Allocation " + str(alloc["id"]) + " for " + node_internal_name
|
|
if alloc["id"] in allocation_server_lookup.keys():
|
|
data = allocation_server_lookup[alloc["id"]]
|
|
svr = data["server"]
|
|
res = data["resources"]
|
|
label = svr["name"] + "(" + svr["identifier"] + ") allocation"
|
|
if res["current_state"] != "running":
|
|
print(f"Excluding allocation \"{label}\" from forwarding, attached server not running")
|
|
continue
|
|
|
|
if "pf_exclude" in svr["description"]:
|
|
print(f"Excluding allocation \"{label}\" from forwarding, found \"pf_exclude\" in the description")
|
|
continue
|
|
|
|
if int(alloc["port"]) <= 1024:
|
|
print(f"Excluding allocation \"{label}\" from forwarding, attempting to allocate port {alloc['port']}, which is deemed a system port")
|
|
continue
|
|
rules.append(create_rule(label, node_internal_ip, alloc["port"]))
|
|
|
|
run(router.async_apply_port_forwarding_rules(rules))
|
|
|
|
print("Currently open ports: ")
|
|
for rule in rules:
|
|
print("\t" + rule_to_string(rule))
|
|
|
|
run(router.async_disconnect())
|
|
run(session.close())
|
|
|
|
|
|
app = Flask(__name__)
|
|
|
|
|
|
@app.route("/pelican-wh", methods=["GET", "POST"])
|
|
def start():
|
|
apply_port_changes()
|
|
return "Hello, world!"
|
|
|
|
|
|
if __name__ == '__main__':
|
|
if os.path.exists("/config/config.ini"):
|
|
cfg.read("/config/config.ini")
|
|
elif os.path.exists("./config.ini"):
|
|
cfg.read("./config.ini")
|
|
else:
|
|
print("No config file detected, exiting...")
|
|
exit(1)
|
|
|
|
port = int(cfg["app"]["port"])
|
|
|
|
print("Applying initial port forwarding")
|
|
apply_port_changes()
|
|
app.run(host="0.0.0.0", port=port)
|