From 823ff884b9324b7af8cc53fd54bf678172b017b1 Mon Sep 17 00:00:00 2001
From: Nick Guy <nick.guy@hotmail.co.uk>
Date: Fri, 7 Mar 2025 18:18:11 +0000
Subject: [PATCH] Update to respect whether the associated server is online

---
 api/__init__.py | 11 ++++++++++-
 api/pelican.py  | 17 ++++++++++++++---
 main.py         | 31 ++++++++++++++++++++++++-------
 3 files changed, 48 insertions(+), 11 deletions(-)

diff --git a/api/__init__.py b/api/__init__.py
index 65b893e..dc31308 100644
--- a/api/__init__.py
+++ b/api/__init__.py
@@ -15,7 +15,7 @@ class Api(object):
     def transform(self, data):
         return data
 
-    def apply_authentication(self) -> Tuple[ApiAuthType, dict]:
+    def apply_authentication(self, use_alt_auth: bool = False) -> Tuple[ApiAuthType, dict]:
         raise NotImplementedError
 
     def _get(self, endpoint, raw: bool = False):
@@ -35,6 +35,15 @@ class Api(object):
             cookies.update(auth)
 
         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:
             return response
 
diff --git a/api/pelican.py b/api/pelican.py
index 6782558..d7f223b 100644
--- a/api/pelican.py
+++ b/api/pelican.py
@@ -3,11 +3,16 @@ from typing import Tuple
 
 
 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.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, {
             "Authorization": "Bearer " + self.auth_token
         }
@@ -17,7 +22,7 @@ class PelicanApi(Api):
         if type == "list":
             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
@@ -30,3 +35,9 @@ class PelicanApi(Api):
 
     def allocations(self, node_id: int):
         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")
diff --git a/main.py b/main.py
index f37d8ac..992f76c 100644
--- a/main.py
+++ b/main.py
@@ -31,7 +31,7 @@ def apply_port_changes():
 
     pelican_cfg = cfg["pelican"]
     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(
         hostname=router_cfg["hostname"],
         username=router_cfg["username"],
@@ -52,8 +52,15 @@ def apply_port_changes():
 
     # Add new auto-generated rules
     allocation_server_lookup = {}
-    for server in pelican.servers():
-        allocation_server_lookup[server["allocation"]] = server
+    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"]
@@ -68,16 +75,24 @@ def apply_port_changes():
 
             label = "Allocation " + str(alloc["id"]) + " for " + node_internal_name
             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"
+                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:
-                rules.append(create_rule(label, node_internal_ip, alloc["port"]))
+            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))
+    # run(router.async_apply_port_forwarding_rules(rules))
 
     print("Currently open ports: ")
     for rule in rules:
@@ -99,6 +114,8 @@ def start():
 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)