352 lines
7.5 KiB
Python
352 lines
7.5 KiB
Python
import os
|
|
import pty
|
|
import fcntl
|
|
import subprocess
|
|
import sqlite3
|
|
import asyncio
|
|
import time
|
|
import re
|
|
from fastapi import FastAPI, WebSocket, BackgroundTasks, Request, Form, WebSocketDisconnect
|
|
from fastapi.responses import RedirectResponse, JSONResponse
|
|
from fastapi.templating import Jinja2Templates
|
|
from fastapi.staticfiles import StaticFiles
|
|
|
|
DB_PATH = "cluster.db"
|
|
SSH_KEY = os.path.expanduser("~/.ssh/id_rsa")
|
|
|
|
app = FastAPI()
|
|
static_path = os.path.join(os.path.dirname(__file__), "static")
|
|
app.mount("/static", StaticFiles(directory=static_path), name="static")
|
|
templates = Jinja2Templates(directory="templates")
|
|
|
|
# -----------------------------
|
|
# DATABASE
|
|
# -----------------------------
|
|
|
|
def init_db():
|
|
conn = sqlite3.connect(DB_PATH)
|
|
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS nodes (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
name TEXT,
|
|
ip TEXT UNIQUE,
|
|
user TEXT,
|
|
|
|
ssh_password TEXT,
|
|
sudo_password TEXT,
|
|
|
|
os TEXT,
|
|
architecture TEXT,
|
|
|
|
docker_installed INTEGER DEFAULT 0,
|
|
vnc_available INTEGER DEFAULT 0,
|
|
|
|
last_seen INTEGER,
|
|
status TEXT
|
|
)
|
|
""")
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
init_db()
|
|
|
|
|
|
def get_db():
|
|
conn = sqlite3.connect(DB_PATH)
|
|
conn.row_factory = sqlite3.Row
|
|
return conn
|
|
|
|
# -----------------------------
|
|
# NODE DISCOVERY
|
|
# -----------------------------
|
|
|
|
async def detect_node_info(ip, user):
|
|
|
|
cmds = {
|
|
"os": "grep '^ID=' /etc/os-release 2>/dev/null | cut -d= -f2",
|
|
"arch": "uname -m",
|
|
"docker": "command -v docker >/dev/null 2>&1 && echo 1 || echo 0",
|
|
"vnc": "pgrep -f vnc >/dev/null 2>&1 && echo 1 || echo 0"
|
|
}
|
|
|
|
results = {}
|
|
|
|
for key, cmd in cmds.items():
|
|
ssh_cmd = f"ssh -o ConnectTimeout=3 -o StrictHostKeyChecking=no {user}@{ip} '{cmd}'"
|
|
|
|
try:
|
|
out = subprocess.check_output(ssh_cmd, shell=True).decode().strip()
|
|
except:
|
|
out = ""
|
|
|
|
results[key] = out
|
|
|
|
return results
|
|
|
|
# -----------------------------
|
|
# NODE MONITOR
|
|
# -----------------------------
|
|
|
|
async def node_monitor_loop():
|
|
|
|
while True:
|
|
|
|
conn = get_db()
|
|
nodes = conn.execute("SELECT * FROM nodes").fetchall()
|
|
|
|
for node in nodes:
|
|
|
|
ip = node["ip"]
|
|
user = node["user"]
|
|
|
|
try:
|
|
|
|
ssh_cmd = f"ssh -o ConnectTimeout=2 -o StrictHostKeyChecking=no {user}@{ip} 'echo online'"
|
|
subprocess.check_output(ssh_cmd, shell=True)
|
|
|
|
conn.execute(
|
|
"UPDATE nodes SET status=?, last_seen=? WHERE id=?",
|
|
("Online", int(time.time()), node["id"])
|
|
)
|
|
|
|
except:
|
|
|
|
conn.execute(
|
|
"UPDATE nodes SET status=? WHERE id=?",
|
|
("Offline", node["id"])
|
|
)
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
await asyncio.sleep(60)
|
|
|
|
|
|
@app.on_event("startup")
|
|
async def start_monitor():
|
|
asyncio.create_task(node_monitor_loop())
|
|
|
|
# -----------------------------
|
|
# SSH BOOTSTRAP
|
|
# -----------------------------
|
|
|
|
async def bootstrap_node(ip, user, password):
|
|
|
|
ssh_copy_cmd = f"sshpass -p '{password}' ssh-copy-id -o StrictHostKeyChecking=no -i {SSH_KEY}.pub {user}@{ip}"
|
|
|
|
subprocess.call(ssh_copy_cmd, shell=True)
|
|
|
|
info = await detect_node_info(ip, user)
|
|
|
|
conn = get_db()
|
|
|
|
conn.execute(
|
|
"""
|
|
UPDATE nodes
|
|
SET os=?, architecture=?, docker_installed=?, vnc_available=?, status=?
|
|
WHERE ip=?
|
|
""",
|
|
(
|
|
info["os"],
|
|
info["arch"],
|
|
int(info["docker"] or 0),
|
|
int(info["vnc"] or 0),
|
|
"Online",
|
|
ip
|
|
)
|
|
)
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
# -----------------------------
|
|
# ROUTES
|
|
# -----------------------------
|
|
|
|
@app.get("/")
|
|
async def index(request: Request):
|
|
|
|
conn = get_db()
|
|
nodes = conn.execute("SELECT * FROM nodes").fetchall()
|
|
conn.close()
|
|
|
|
return templates.TemplateResponse("index.html", {"request": request, "nodes": nodes})
|
|
|
|
|
|
@app.post("/add_node")
|
|
async def add_node(
|
|
background_tasks: BackgroundTasks,
|
|
name: str = Form(...),
|
|
ip: str = Form(...),
|
|
user: str = Form(...),
|
|
password: str = Form(...)
|
|
):
|
|
|
|
conn = get_db()
|
|
|
|
conn.execute(
|
|
"""
|
|
INSERT INTO nodes (name, ip, user, ssh_password, status)
|
|
VALUES (?, ?, ?, ?, ?)
|
|
""",
|
|
(name, ip, user, password, "Connecting")
|
|
)
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
background_tasks.add_task(bootstrap_node, ip, user, password)
|
|
|
|
return RedirectResponse("/", 303)
|
|
|
|
|
|
@app.post("/remove_node/{node_id}")
|
|
async def remove_node(node_id: int):
|
|
|
|
conn = get_db()
|
|
conn.execute("DELETE FROM nodes WHERE id=?", (node_id,))
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
return RedirectResponse("/", 303)
|
|
|
|
|
|
@app.get("/refresh_status/{node_id}")
|
|
async def refresh_status(node_id: int):
|
|
|
|
conn = get_db()
|
|
node = conn.execute("SELECT * FROM nodes WHERE id=?", (node_id,)).fetchone()
|
|
|
|
if not node:
|
|
return JSONResponse({"status": "Unknown"})
|
|
|
|
ip = node["ip"]
|
|
user = node["user"]
|
|
|
|
try:
|
|
|
|
ssh_cmd = f"ssh -o ConnectTimeout=2 -o StrictHostKeyChecking=no {user}@{ip} 'echo online'"
|
|
subprocess.check_output(ssh_cmd, shell=True)
|
|
|
|
status = "Online"
|
|
|
|
except:
|
|
|
|
status = "Offline"
|
|
|
|
conn.execute("UPDATE nodes SET status=? WHERE id=?", (status, node_id))
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
return {"status": status}
|
|
|
|
# -----------------------------
|
|
# NODE EDIT
|
|
# -----------------------------
|
|
|
|
@app.get("/node/{node_id}")
|
|
async def get_node(node_id: int):
|
|
|
|
conn = get_db()
|
|
node = conn.execute("SELECT * FROM nodes WHERE id=?", (node_id,)).fetchone()
|
|
conn.close()
|
|
|
|
if not node:
|
|
return JSONResponse({})
|
|
|
|
return dict(node)
|
|
|
|
|
|
@app.post("/update_node/{node_id}")
|
|
async def update_node(
|
|
node_id: int,
|
|
name: str = Form(...),
|
|
ip: str = Form(...),
|
|
user: str = Form(...)
|
|
):
|
|
|
|
conn = get_db()
|
|
|
|
conn.execute(
|
|
"""
|
|
UPDATE nodes
|
|
SET name=?, ip=?, user=?
|
|
WHERE id=?
|
|
""",
|
|
(name, ip, user, node_id)
|
|
)
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
return RedirectResponse("/", 303)
|
|
|
|
# -----------------------------
|
|
# TERMINAL
|
|
# -----------------------------
|
|
|
|
@app.websocket("/ws/terminal/{ip}")
|
|
async def terminal(websocket: WebSocket, ip: str):
|
|
|
|
await websocket.accept()
|
|
|
|
conn = get_db()
|
|
node = conn.execute("SELECT * FROM nodes WHERE ip=?", (ip,)).fetchone()
|
|
conn.close()
|
|
|
|
if not node:
|
|
await websocket.close()
|
|
return
|
|
|
|
master_fd, slave_fd = pty.openpty()
|
|
|
|
proc = await asyncio.create_subprocess_exec(
|
|
"ssh",
|
|
"-o",
|
|
"StrictHostKeyChecking=no",
|
|
"-t",
|
|
f"{node['user']}@{ip}",
|
|
stdin=slave_fd,
|
|
stdout=slave_fd,
|
|
stderr=slave_fd
|
|
)
|
|
|
|
async def pty_to_ws():
|
|
|
|
fl = fcntl.fcntl(master_fd, fcntl.F_GETFL)
|
|
fcntl.fcntl(master_fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
|
|
|
|
while True:
|
|
|
|
await asyncio.sleep(0.01)
|
|
|
|
try:
|
|
data = os.read(master_fd, 1024).decode(errors="ignore")
|
|
|
|
if data:
|
|
await websocket.send_text(data)
|
|
|
|
except:
|
|
pass
|
|
|
|
async def ws_to_pty():
|
|
|
|
try:
|
|
while True:
|
|
|
|
data = await websocket.receive_text()
|
|
os.write(master_fd, data.encode())
|
|
|
|
except:
|
|
pass
|
|
|
|
await asyncio.gather(pty_to_ws(), ws_to_pty())
|
|
|
|
if __name__ == "__main__":
|
|
|
|
import uvicorn
|
|
|
|
uvicorn.run(app, host="0.0.0.0", port=8000)
|