Added suport for API keys on cluster level

main.py

@@ -12,10 +12,15 @@ from fastapi.openapi.utils import get_openapi
 def init_db():
     conn = sqlite3.connect("servers.db")
     cursor = conn.cursor()
+    cursor.execute('''CREATE TABLE IF NOT EXISTS clusters (
+                        name TEXT PRIMARY KEY, 
+                        api_token TEXT NOT NULL)''')
     cursor.execute('''CREATE TABLE IF NOT EXISTS servers (
                         name TEXT PRIMARY KEY, 
                         ip TEXT NOT NULL, 
-                        mac TEXT NOT NULL)''')
+                        mac TEXT NOT NULL, 
+                        cluster_name TEXT NOT NULL,
+                        FOREIGN KEY(cluster_name) REFERENCES clusters(name))''')
     conn.commit()
     conn.close()
 
@@ -24,25 +29,40 @@ init_db()
 def get_server_info(server_name):
     conn = sqlite3.connect("servers.db")
     cursor = conn.cursor()
-    cursor.execute("SELECT ip, mac FROM servers WHERE name = ?", (server_name,))
+    cursor.execute("SELECT servers.ip, servers.mac, clusters.api_token FROM servers JOIN clusters ON servers.cluster_name = clusters.name WHERE servers.name = ?", (server_name,))
     server = cursor.fetchone()
     conn.close()
     if server:
-        return {"ip": server[0], "mac": server[1]}
+        return {"ip": server[0], "mac": server[1], "api_token": server[2]}
     return None
 
 def list_servers():
     conn = sqlite3.connect("servers.db")
     cursor = conn.cursor()
-    cursor.execute("SELECT name, ip, mac FROM servers")
+    cursor.execute("SELECT servers.name, servers.ip, servers.mac, servers.cluster_name FROM servers")
     servers = cursor.fetchall()
     conn.close()
-    return [{"name": s[0], "ip": s[1], "mac": s[2]} for s in servers]
+    return [{"name": s[0], "ip": s[1], "mac": s[2], "cluster_name": s[3]} for s in servers]
 
-def add_or_update_server(name: str, ip: str, mac: str):
+def list_clusters():
     conn = sqlite3.connect("servers.db")
     cursor = conn.cursor()
-    cursor.execute("REPLACE INTO servers (name, ip, mac) VALUES (?, ?, ?)", (name, ip, mac))
+    cursor.execute("SELECT name FROM clusters")
+    clusters = cursor.fetchall()
+    conn.close()
+    return [c[0] for c in clusters]
+
+def add_or_update_cluster(name: str, api_token: str):
+    conn = sqlite3.connect("servers.db")
+    cursor = conn.cursor()
+    cursor.execute("REPLACE INTO clusters (name, api_token) VALUES (?, ?)", (name, api_token))
+    conn.commit()
+    conn.close()
+
+def add_or_update_server(name: str, ip: str, mac: str, cluster_name: str):
+    conn = sqlite3.connect("servers.db")
+    cursor = conn.cursor()
+    cursor.execute("REPLACE INTO servers (name, ip, mac, cluster_name) VALUES (?, ?, ?, ?)", (name, ip, mac, cluster_name))
     conn.commit()
     conn.close()
 
@@ -52,11 +72,17 @@ class ServerModel(BaseModel):
     name: str
     ip: str
     mac: str
+    cluster_name: str
 
-def get_proxmox_status(server_ip):
-    """Fetch Proxmox server status."""
+class ClusterModel(BaseModel):
+    name: str
+    api_token: str
+
+def get_proxmox_status(server_ip, api_token):
+    """Fetch Proxmox server status using API token."""
+    headers = {"Authorization": f"PVEAPIToken={api_token}"}
     try:
-        response = requests.get(f"https://{server_ip}:8006/api2/json/nodes", verify=False)  # Assuming no authentication
+        response = requests.get(f"https://{server_ip}:8006/api2/json/nodes", headers=headers, verify=False)
         response.raise_for_status()
         data = response.json()
         return {"status": data["data"][0]["status"]}  # Extract node status
@@ -90,13 +116,28 @@ def status(server_name: str):
     server = get_server_info(server_name)
     if not server:
         raise HTTPException(status_code=404, detail="Server not found")
-    return get_proxmox_status(server["ip"])
+    return get_proxmox_status(server["ip"], server["api_token"])
 
 @app.get("/statuses")
 def list_all_statuses():
     """Returns the statuses of all servers."""
     servers = list_servers()
-    return {server["name"]: get_proxmox_status(server["ip"]) for server in servers}
+    statuses = {}
+    for server in servers:
+        cluster_token = get_server_info(server["name"])["api_token"]
+        statuses[server["name"]] = get_proxmox_status(server["ip"], cluster_token)
+    return statuses
+
+@app.get("/clusters")
+def get_clusters():
+    """Returns a list of all cluster names."""
+    return list_clusters()
+
+@app.post("/clusters")
+def add_cluster(cluster: ClusterModel):
+    """Adds or updates a cluster."""
+    add_or_update_cluster(cluster.name, cluster.api_token)
+    return {"message": "Cluster added/updated successfully"}
 
 def check_power_state(server_ip):
     """Check if the server is online by pinging it."""
@@ -147,7 +188,7 @@ def get_servers():
 @app.post("/servers")
 def add_server(server: ServerModel):
     """Adds or updates a server."""
-    add_or_update_server(server.name, server.ip, server.mac)
+    add_or_update_server(server.name, server.ip, server.mac, server.cluster_name)
     return {"message": "Server added/updated successfully"}
 
-# Run the server with: uvicorn script_name:app --host 0.0.0.0 --port 8000
+# Run the server with: uvicorn main:app --host 0.0.0.0 --port 8000