implement user authentication, role-based access control, and multi-user instance isolation

main.py

@@ -11,9 +11,12 @@ from fastapi import FastAPI, Request, Depends, HTTPException
 from fastapi.responses import Response, FileResponse, JSONResponse
 from fastapi.staticfiles import StaticFiles
 from starlette.middleware.sessions import SessionMiddleware
+from passlib.context import CryptContext
 from pydantic import BaseModel
 from cryptography.fernet import Fernet
 
+pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
+
 WEB_PASSWORD = os.getenv("WEB_PASSWORD", "admin")
 ENCRYPTION_KEY = os.getenv("ENCRYPTION_KEY")
 FIREWALL_HOST_IP = os.getenv("FIREWALL_HOST_IP")
@@ -33,9 +36,35 @@ except ValueError:
 app = FastAPI()
 app.add_middleware(SessionMiddleware, secret_key=ENCRYPTION_KEY)
 
+def get_password_hash(password):
+    return pwd_context.hash(password)
+
+def verify_password(plain_password, hashed_password):
+    return pwd_context.verify(plain_password, hashed_password)
+
 def init_db():
     conn = sqlite3.connect(DB_PATH)
     c = conn.cursor()
+    
+    # Create users table
+    c.execute('''
+        CREATE TABLE IF NOT EXISTS users (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            username TEXT UNIQUE NOT NULL,
+            hashed_password TEXT NOT NULL,
+            is_admin INTEGER DEFAULT 0
+        )
+    ''')
+    
+    # Check if instances table has user_id, if not, add it (simple migration)
+    c.execute("PRAGMA table_info(instances)")
+    columns = [col[1] for col in c.fetchall()]
+    if "user_id" not in columns:
+        print("Migrating instances table to include user_id...")
+        # SQLite doesn't support adding a column with a foreign key constraint easily,
+        # so we'll just add the column for now.
+        c.execute('ALTER TABLE instances ADD COLUMN user_id INTEGER')
+    
     c.execute('''
         CREATE TABLE IF NOT EXISTS instances (
             id INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -44,27 +73,58 @@ def init_db():
             port INTEGER NOT NULL,
             encrypted_secret TEXT NOT NULL,
             last_status TEXT DEFAULT 'Unknown',
-            last_checked REAL DEFAULT 0
+            last_checked REAL DEFAULT 0,
+            user_id INTEGER,
+            FOREIGN KEY (user_id) REFERENCES users (id)
         )
     ''')
+    
+    # Create default admin if no users exist
+    c.execute('SELECT COUNT(*) FROM users')
+    if c.fetchone()[0] == 0:
+        print("Creating default admin user...")
+        admin_hash = get_password_hash(WEB_PASSWORD)
+        c.execute('INSERT INTO users (username, hashed_password, is_admin) VALUES (?, ?, ?)',
+                  ('admin', admin_hash, 1))
+        
+        # Associate existing instances with the new admin user
+        admin_id = c.lastrowid
+        c.execute('UPDATE instances SET user_id = ? WHERE user_id IS NULL', (admin_id,))
+        
     conn.commit()
     conn.close()
 
 init_db()
 
 def is_authenticated(request: Request):
-    if not request.session.get("authenticated"):
+    if not request.session.get("user_id"):
         raise HTTPException(status_code=401, detail="Not authenticated")
     return True
 
+def is_admin(request: Request):
+    if not request.session.get("is_admin"):
+        raise HTTPException(status_code=403, detail="Admin access required")
+    return True
+
 @app.post("/api/login")
 async def login(request: Request):
     data = await request.json()
+    username = data.get("username")
     password = data.get("password")
-    if secrets.compare_digest(password, WEB_PASSWORD):
-        request.session["authenticated"] = True
-        return {"status": "ok"}
-    raise HTTPException(status_code=401, detail="Invalid password")
+    
+    conn = sqlite3.connect(DB_PATH)
+    conn.row_factory = sqlite3.Row
+    c = conn.cursor()
+    c.execute('SELECT * FROM users WHERE username = ?', (username,))
+    user = c.fetchone()
+    conn.close()
+    
+    if user and verify_password(password, user['hashed_password']):
+        request.session["user_id"] = user['id']
+        request.session["username"] = user['username']
+        request.session["is_admin"] = bool(user['is_admin'])
+        return {"status": "ok", "user": {"username": user['username'], "is_admin": bool(user['is_admin'])}}
+    raise HTTPException(status_code=401, detail="Invalid username or password")
 
 @app.get("/api/logout")
 async def logout(request: Request):
@@ -73,7 +133,91 @@ async def logout(request: Request):
 
 @app.get("/api/auth/status")
 async def auth_status(request: Request):
-    return {"authenticated": request.session.get("authenticated", False)}
+    if request.session.get("user_id"):
+        return {
+            "authenticated": True,
+            "user": {
+                "username": request.session.get("username"),
+                "is_admin": request.session.get("is_admin")
+            }
+        }
+    return {"authenticated": False}
+
+# User Management (Admin Only)
+@app.get("/api/users", dependencies=[Depends(is_authenticated), Depends(is_admin)])
+def get_users():
+    conn = sqlite3.connect(DB_PATH)
+    conn.row_factory = sqlite3.Row
+    c = conn.cursor()
+    c.execute('SELECT id, username, is_admin FROM users')
+    users = [dict(row) for row in c.fetchall()]
+    conn.close()
+    return users
+
+@app.post("/api/users", dependencies=[Depends(is_authenticated), Depends(is_admin)])
+async def create_user(request: Request):
+    data = await request.json()
+    username = data.get("username")
+    password = data.get("password")
+    is_admin = 1 if data.get("is_admin") else 0
+    
+    if not username or not password:
+        raise HTTPException(status_code=400, detail="Username and password required")
+        
+    hashed = get_password_hash(password)
+    try:
+        conn = sqlite3.connect(DB_PATH)
+        c = conn.cursor()
+        c.execute('INSERT INTO users (username, hashed_password, is_admin) VALUES (?, ?, ?)',
+                  (username, hashed, is_admin))
+        conn.commit()
+        conn.close()
+    except sqlite3.IntegrityError:
+        raise HTTPException(status_code=400, detail="Username already exists")
+    return {"status": "ok"}
+
+@app.delete("/api/users/{id}", dependencies=[Depends(is_authenticated), Depends(is_admin)])
+def delete_user(id: int, request: Request):
+    if id == request.session.get("user_id"):
+        raise HTTPException(status_code=400, detail="Cannot delete yourself")
+    conn = sqlite3.connect(DB_PATH)
+    c = conn.cursor()
+    c.execute('DELETE FROM users WHERE id=?', (id,))
+    # Also delete their instances? Let's say yes for cleanliness.
+    c.execute('DELETE FROM instances WHERE user_id=?', (id,))
+    conn.commit()
+    conn.close()
+    return {"status": "ok"}
+
+@app.put("/api/users/{id}/password", dependencies=[Depends(is_authenticated), Depends(is_admin)])
+async def admin_reset_password(id: int, request: Request):
+    data = await request.json()
+    new_password = data.get("password")
+    if not new_password:
+        raise HTTPException(status_code=400, detail="New password required")
+    hashed = get_password_hash(new_password)
+    conn = sqlite3.connect(DB_PATH)
+    c = conn.cursor()
+    c.execute('UPDATE users SET hashed_password = ? WHERE id = ?', (hashed, id))
+    conn.commit()
+    conn.close()
+    return {"status": "ok"}
+
+# Own Password Change (All Users)
+@app.put("/api/users/me/password", dependencies=[Depends(is_authenticated)])
+async def change_own_password(request: Request):
+    data = await request.json()
+    new_password = data.get("password")
+    if not new_password:
+        raise HTTPException(status_code=400, detail="New password required")
+    user_id = request.session.get("user_id")
+    hashed = get_password_hash(new_password)
+    conn = sqlite3.connect(DB_PATH)
+    c = conn.cursor()
+    c.execute('UPDATE users SET hashed_password = ? WHERE id = ?', (hashed, user_id))
+    conn.commit()
+    conn.close()
+    return {"status": "ok"}
 
 class InstanceCreate(BaseModel):
     name: str
@@ -153,25 +297,27 @@ async def startup_event():
     asyncio.create_task(poll_instances())
 
 @app.get("/api/instances", dependencies=[Depends(is_authenticated)])
-def get_instances():
+def get_instances(request: Request):
+    user_id = request.session.get("user_id")
     conn = sqlite3.connect(DB_PATH)
     conn.row_factory = sqlite3.Row
     c = conn.cursor()
-    c.execute('SELECT id, name, ip, port, last_status, last_checked FROM instances')
+    c.execute('SELECT id, name, ip, port, last_status, last_checked FROM instances WHERE user_id = ?', (user_id,))
     instances = [dict(row) for row in c.fetchall()]
     conn.close()
     return instances
 
 @app.post("/api/instances", dependencies=[Depends(is_authenticated)])
-def create_instance(inst: InstanceCreate):
+def create_instance(inst: InstanceCreate, request: Request):
+    user_id = request.session.get("user_id")
     cleaned = clean_secret(inst.secret)
     encrypted = fernet.encrypt(cleaned.encode()).decode()
     conn = sqlite3.connect(DB_PATH)
     c = conn.cursor()
     c.execute('''
-        INSERT INTO instances (name, ip, port, encrypted_secret)
-        VALUES (?, ?, ?, ?)
-    ''', (inst.name, inst.ip, inst.port, encrypted))
+        INSERT INTO instances (name, ip, port, encrypted_secret, user_id)
+        VALUES (?, ?, ?, ?, ?)
+    ''', (inst.name, inst.ip, inst.port, encrypted, user_id))
     conn.commit()
     conn.close()
     return {"status": "ok"}
@@ -181,10 +327,11 @@ def get_config():
     return {"firewall_host_ip": FIREWALL_HOST_IP}
 
 @app.delete("/api/instances/{id}", dependencies=[Depends(is_authenticated)])
-def delete_instance(id: int):
+def delete_instance(id: int, request: Request):
+    user_id = request.session.get("user_id")
     conn = sqlite3.connect(DB_PATH)
     c = conn.cursor()
-    c.execute('DELETE FROM instances WHERE id=?', (id,))
+    c.execute('DELETE FROM instances WHERE id=? AND user_id=?', (id, user_id))
     conn.commit()
     conn.close()
     return {"status": "ok"}