diff --git a/main.py b/main.py index fc81fce..7ebe1ed 100644 --- a/main.py +++ b/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 diff --git a/servers.db b/servers.db index 6fe74ce..a6fa9ed 100644 Binary files a/servers.db and b/servers.db differ