Update to respect whether the associated server is online

This commit is contained in:
Nick Guy 2025-03-07 18:18:11 +00:00
parent ca7e03bd66
commit 823ff884b9
3 changed files with 48 additions and 11 deletions

View file

@ -15,7 +15,7 @@ class Api(object):
def transform(self, data): def transform(self, data):
return data return data
def apply_authentication(self) -> Tuple[ApiAuthType, dict]: def apply_authentication(self, use_alt_auth: bool = False) -> Tuple[ApiAuthType, dict]:
raise NotImplementedError raise NotImplementedError
def _get(self, endpoint, raw: bool = False): def _get(self, endpoint, raw: bool = False):
@ -35,6 +35,15 @@ class Api(object):
cookies.update(auth) cookies.update(auth)
response = requests.get(url, headers=headers, cookies=cookies) response = requests.get(url, headers=headers, cookies=cookies)
if response.status_code == 403:
ty, auth = self.apply_authentication(True)
if ty == ApiAuthType.Header:
headers.update(auth)
elif ty == ApiAuthType.Cookie:
cookies.update(auth)
response = requests.get(url, headers=headers, cookies=cookies)
if raw: if raw:
return response return response

View file

@ -3,11 +3,16 @@ from typing import Tuple
class PelicanApi(Api): class PelicanApi(Api):
def __init__(self, base_url, auth_token): def __init__(self, base_url, auth_token, client_auth_token):
self.base_url = base_url self.base_url = base_url
self.auth_token = auth_token self.auth_token = auth_token
self.client_auth_token = client_auth_token
def apply_authentication(self) -> Tuple[ApiAuthType, dict]: def apply_authentication(self, use_alt_auth: bool = False) -> Tuple[ApiAuthType, dict]:
if use_alt_auth:
return ApiAuthType.Header, {
"Authorization": "Bearer " + self.client_auth_token
}
return ApiAuthType.Header, { return ApiAuthType.Header, {
"Authorization": "Bearer " + self.auth_token "Authorization": "Bearer " + self.auth_token
} }
@ -17,7 +22,7 @@ class PelicanApi(Api):
if type == "list": if type == "list":
return [self.transform(x) for x in data["data"]] return [self.transform(x) for x in data["data"]]
if type in ["allocation", "server", "node"]: if type in ["allocation", "server", "node", "stats"]:
return data["attributes"] return data["attributes"]
return data return data
@ -30,3 +35,9 @@ class PelicanApi(Api):
def allocations(self, node_id: int): def allocations(self, node_id: int):
return self._get("/application/nodes/" + str(node_id) + "/allocations") return self._get("/application/nodes/" + str(node_id) + "/allocations")
def server_allocations(self, identifier):
return self._get("/client/servers/" + str(identifier) + "/network/allocations")
def resources(self, identifier):
return self._get("/client/servers/" + str(identifier) + "/resources")

31
main.py
View file

@ -31,7 +31,7 @@ def apply_port_changes():
pelican_cfg = cfg["pelican"] pelican_cfg = cfg["pelican"]
router_cfg = cfg["router"] router_cfg = cfg["router"]
pelican = PelicanApi(pelican_cfg["baseurl"], pelican_cfg["auth"]) pelican = PelicanApi(pelican_cfg["baseurl"], pelican_cfg["auth"], pelican_cfg["client_auth"])
router = AsusRouter( router = AsusRouter(
hostname=router_cfg["hostname"], hostname=router_cfg["hostname"],
username=router_cfg["username"], username=router_cfg["username"],
@ -52,8 +52,15 @@ def apply_port_changes():
# Add new auto-generated rules # Add new auto-generated rules
allocation_server_lookup = {} allocation_server_lookup = {}
for server in pelican.servers(): servers = pelican.servers()
allocation_server_lookup[server["allocation"]] = server 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(): for node in pelican.nodes():
node_id = node["id"] node_id = node["id"]
@ -68,16 +75,24 @@ def apply_port_changes():
label = "Allocation " + str(alloc["id"]) + " for " + node_internal_name label = "Allocation " + str(alloc["id"]) + " for " + node_internal_name
if alloc["id"] in allocation_server_lookup.keys(): if alloc["id"] in allocation_server_lookup.keys():
svr = allocation_server_lookup[alloc["id"]] data = allocation_server_lookup[alloc["id"]]
svr = data["server"]
res = data["resources"]
label = svr["name"] + "(" + svr["identifier"] + ") allocation" 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"]: if "pf_exclude" in svr["description"]:
print(f"Excluding allocation \"{label}\" from forwarding, found \"pf_exclude\" in the description") print(f"Excluding allocation \"{label}\" from forwarding, found \"pf_exclude\" in the description")
continue continue
if int(alloc["port"]) > 1024: if int(alloc["port"]) <= 1024:
rules.append(create_rule(label, node_internal_ip, alloc["port"])) 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)) # run(router.async_apply_port_forwarding_rules(rules))
print("Currently open ports: ") print("Currently open ports: ")
for rule in rules: for rule in rules:
@ -99,6 +114,8 @@ def start():
if __name__ == '__main__': if __name__ == '__main__':
if os.path.exists("/config/config.ini"): if os.path.exists("/config/config.ini"):
cfg.read("/config/config.ini") cfg.read("/config/config.ini")
elif os.path.exists("./config.ini"):
cfg.read("./config.ini")
else: else:
print("No config file detected, exiting...") print("No config file detected, exiting...")
exit(1) exit(1)