feat:

--Added JWT Token session

main.py

@@ -1,10 +1,28 @@
-from fastapi import FastAPI, HTTPException,WebSocket, WebSocketDisconnect
+from fastapi import FastAPI, HTTPException,WebSocket, WebSocketDisconnect, Header,Depends,Query,status
+from datetime import datetime, timedelta, timezone
+import jwt
+from jwt import ExpiredSignatureError, InvalidTokenError
 from pydantic import BaseModel
 import psycopg2
 from psycopg2.extras import RealDictCursor
 import asyncio
+from fastapi.middleware.cors import CORSMiddleware
+origins = [
+    "http://localhost:5173",
+    "http://127.0.0.1:5173",
+]
+
+
 app = FastAPI()
 
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=origins,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
 
 class RegisterNodeRequest(BaseModel):
     node_name: str
@@ -48,9 +66,97 @@ class LinkNodeRequest(BaseModel):
 # );
 
 
+# ============================================================
+# JWT TOKEN CREATION
+# ============================================================
+SECRET_KEY = "my-super-secret-key-change-this"
+ALGORITHM = "HS256"
 
+
+@app.post("/get-token")
+def get_token():
+    print("hio")
+    payload = {
+        "sub": "frontend",
+        "exp": datetime.now(timezone.utc) + timedelta(hours=1)
+    }
+
+    token = jwt.encode(
+        payload,
+        SECRET_KEY,
+        algorithm=ALGORITHM
+    )
+
+    return {
+        "access_token": token,
+        "token_type": "bearer"
+    }
+
+
+
+
+# ============================================================
+# JWT VERIFICATION
+# ============================================================
+def verify_token(
+    authorization: str = Header(None),
+):
+
+    if not authorization:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Authorization header missing",
+            headers={
+                "WWW-Authenticate": "Bearer"
+            },
+        )
+
+    parts = authorization.split()
+
+    if len(parts) != 2:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid Authorization header",
+        )
+
+    scheme = parts[0]
+    token = parts[1]
+
+    if scheme.lower() != "bearer":
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Authorization scheme must be Bearer",
+        )
+
+    try:
+
+        payload = jwt.decode(
+            token,
+            SECRET_KEY,
+            algorithms=[ALGORITHM],
+        )
+
+        return payload
+
+    except ExpiredSignatureError:
+
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Token expired",
+        )
+
+    except InvalidTokenError:
+
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid token",
+        )
+
+
+
+# Register Nodes 
 @app.post("/register-node")
-def register_node(data: RegisterNodeRequest):
+def register_node(data: RegisterNodeRequest, token_data: dict = Depends(verify_token)):
 
     conn = None
     cur = None
@@ -126,7 +232,8 @@ def register_node(data: RegisterNodeRequest):
             "message": "Node registered",
             "node_id": node_id,
             "spock_version": spock[1],
-            "nodes": nodes
+            "nodes": nodes,
+            "authenticated_as": token_data["sub"]
         }
 
     except Exception as e:
@@ -158,7 +265,7 @@ def register_node(data: RegisterNodeRequest):
 #   "password": "postgres"
 # }
 @app.post("/link-nodes")
-def link_nodes(data: LinkNodeRequest):
+def link_nodes(data: LinkNodeRequest,token_data: dict = Depends(verify_token)):
 
     conn = None
     cur = None
@@ -215,7 +322,9 @@ def link_nodes(data: LinkNodeRequest):
             "subscription_id": subscription_id,
             "subscription_name": data.subscription_name,
             "subscriber": data.subscriber_host,
-            "provider": data.provider_host
+            "provider": data.provider_host,
+            "authenticated_as": token_data["sub"],
+
         }
 
     except Exception as e:
@@ -273,7 +382,18 @@ def get_spock_status():
 
 
 @app.websocket("/ws/spock-status")
-async def websocket_spock_status(websocket: WebSocket):
+async def websocket_spock_status(websocket: WebSocket, token: str = Query(...),):
+    token_data = verify_websocket_token(token)
+
+    if token_data is None:
+
+        await websocket.close(
+            code=1008,
+            reason="Invalid or expired JWT",
+        )
+
+        return
+
     await websocket.accept()
 
     try:
@@ -299,13 +419,15 @@ async def websocket_spock_status(websocket: WebSocket):
             await asyncio.sleep(2)
 
     except WebSocketDisconnect:
-        print("WebSocket client disconnected")
+        print("WebSocket client disconnected" f"{token_data.get('sub')}")
         return True
 
 
 
 @app.get("/spock/subscriptions")
-def get_spock_subscription_status():
+def get_spock_subscription_status(token_data: dict = Depends(verify_token),):
+    conn = None
+    cursor = None
     try:
         conn = psycopg2.connect(**DB_CONFIG)
 
@@ -324,7 +446,9 @@ def get_spock_subscription_status():
         return {
             "success": True,
             "count": len(rows),
-            "subscriptions": rows
+            "subscriptions": rows,
+            "authenticated_as": token_data["sub"],
+
         }
 
     except Exception as e:
@@ -334,3 +458,5 @@ def get_spock_subscription_status():
         )     
 
 
+
+

main_experimental.py

@@ -0,0 +1,507 @@
+from fastapi import (
+    FastAPI,
+    HTTPException,
+    WebSocket,
+    WebSocketDisconnect,
+    Depends,
+    Header,
+    Query,
+    status,
+)
+from datetime import datetime, timedelta, timezone
+from pydantic import BaseModel
+import jwt
+from jwt import ExpiredSignatureError, InvalidTokenError
+import psycopg2
+from psycopg2.extras import RealDictCursor
+import asyncio
+
+
+app = FastAPI()
+
+
+# ============================================================
+# JWT CONFIG
+# ============================================================
+
+SECRET_KEY = "my-super-secret-key-change-this"
+ALGORITHM = "HS256"
+ACCESS_TOKEN_EXPIRE_MINUTES = 60
+
+
+# ============================================================
+# REQUEST MODELS
+# ============================================================
+
+class RegisterNodeRequest(BaseModel):
+    node_name: str
+    host: str
+    port: int = 5444
+    dbname: str = "postgres"
+    user: str = "postgres"
+    password: str
+
+
+class LinkNodeRequest(BaseModel):
+    subscription_name: str
+
+    subscriber_host: str
+    subscriber_port: int = 5444
+
+    provider_host: str
+    provider_port: int = 5444
+
+    dbname: str = "postgres"
+    user: str = "postgres"
+    password: str
+
+
+# ============================================================
+# JWT TOKEN CREATION
+# ============================================================
+
+@app.post("/get-token")
+def get_token():
+
+    payload = {
+        "sub": "frontend",
+        "iat": datetime.now(timezone.utc),
+        "exp": datetime.now(timezone.utc)
+        + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES),
+    }
+
+    token = jwt.encode(
+        payload,
+        SECRET_KEY,
+        algorithm=ALGORITHM,
+    )
+
+    return {
+        "access_token": token,
+        "token_type": "bearer",
+        "expires_in": ACCESS_TOKEN_EXPIRE_MINUTES * 60,
+    }
+
+
+# ============================================================
+# JWT VERIFICATION
+# ============================================================
+
+def verify_token(
+    authorization: str = Header(None),
+):
+
+    if not authorization:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Authorization header missing",
+            headers={
+                "WWW-Authenticate": "Bearer"
+            },
+        )
+
+    parts = authorization.split()
+
+    if len(parts) != 2:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid Authorization header",
+        )
+
+    scheme = parts[0]
+    token = parts[1]
+
+    if scheme.lower() != "bearer":
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Authorization scheme must be Bearer",
+        )
+
+    try:
+
+        payload = jwt.decode(
+            token,
+            SECRET_KEY,
+            algorithms=[ALGORITHM],
+        )
+
+        return payload
+
+    except ExpiredSignatureError:
+
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Token expired",
+        )
+
+    except InvalidTokenError:
+
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid token",
+        )
+
+
+# ============================================================
+# REGISTER NODE
+# JWT PROTECTED
+# ============================================================
+
+@app.post("/register-node")
+def register_node(
+    data: RegisterNodeRequest,
+    token_data: dict = Depends(verify_token),
+):
+
+    conn = None
+    cur = None
+
+    try:
+
+        conn = psycopg2.connect(
+            host=data.host,
+            port=data.port,
+            dbname=data.dbname,
+            user=data.user,
+            password=data.password,
+        )
+
+        cur = conn.cursor()
+
+        # Create Spock extension
+        cur.execute("""
+            CREATE EXTENSION IF NOT EXISTS spock;
+        """)
+
+        # Check version
+        cur.execute("""
+            SELECT extname, extversion
+            FROM pg_extension
+            WHERE extname = 'spock';
+        """)
+
+        spock = cur.fetchone()
+
+        dsn = (
+            f"host={data.host} "
+            f"port={data.port} "
+            f"dbname={data.dbname} "
+            f"user={data.user} "
+            f"password={data.password}"
+        )
+
+        # Register Spock node
+        cur.execute("""
+            SELECT spock.node_create(
+                node_name := %s,
+                dsn := %s
+            );
+        """, (
+            data.node_name,
+            dsn,
+        ))
+
+        node_id = cur.fetchone()[0]
+
+        conn.commit()
+
+        cur.execute("""
+            SELECT node_id, node_name
+            FROM spock.node;
+        """)
+
+        rows = cur.fetchall()
+
+        nodes = [
+            {
+                "node_id": row[0],
+                "node_name": row[1],
+            }
+            for row in rows
+        ]
+
+        return {
+            "status": "success",
+            "message": "Node registered",
+            "node_id": node_id,
+            "spock_version": spock[1],
+            "nodes": nodes,
+            "authenticated_as": token_data["sub"],
+        }
+
+    except Exception as e:
+
+        if conn:
+            conn.rollback()
+
+        raise HTTPException(
+            status_code=500,
+            detail=str(e),
+        )
+
+    finally:
+
+        if cur:
+            cur.close()
+
+        if conn:
+            conn.close()
+
+
+# ============================================================
+# LINK NODES
+# JWT PROTECTED
+# ============================================================
+
+@app.post("/link-nodes")
+def link_nodes(
+    data: LinkNodeRequest,
+    token_data: dict = Depends(verify_token),
+):
+
+    conn = None
+    cur = None
+
+    try:
+
+        conn = psycopg2.connect(
+            host=data.subscriber_host,
+            port=data.subscriber_port,
+            dbname=data.dbname,
+            user=data.user,
+            password=data.password,
+        )
+
+        cur = conn.cursor()
+
+        provider_dsn = (
+            f"host={data.provider_host} "
+            f"port={data.provider_port} "
+            f"dbname={data.dbname} "
+            f"user={data.user} "
+            f"password={data.password}"
+        )
+
+        cur.execute("""
+            SELECT spock.sub_create(
+                subscription_name := %s,
+                provider_dsn := %s,
+                replication_sets := ARRAY[
+                    'default',
+                    'default_insert_only',
+                    'ddl_sql'
+                ],
+                synchronize_structure := false,
+                synchronize_data := false
+            );
+        """, (
+            data.subscription_name,
+            provider_dsn,
+        ))
+
+        subscription_id = cur.fetchone()[0]
+
+        conn.commit()
+
+        return {
+            "status": "success",
+            "message": "Nodes linked successfully",
+            "subscription_id": subscription_id,
+            "subscription_name": data.subscription_name,
+            "subscriber": data.subscriber_host,
+            "provider": data.provider_host,
+            "authenticated_as": token_data["sub"],
+        }
+
+    except Exception as e:
+
+        if conn:
+            conn.rollback()
+
+        raise HTTPException(
+            status_code=500,
+            detail=str(e),
+        )
+
+    finally:
+
+        if cur:
+            cur.close()
+
+        if conn:
+            conn.close()
+
+
+# ============================================================
+# DATABASE CONFIG
+# ============================================================
+
+DB_CONFIG = {
+    "host": "172.16.10.123",
+    "port": 5444,
+    "dbname": "postgres",
+    "user": "postgres",
+    "password": "postgres",
+}
+
+
+# ============================================================
+# SPOCK STATUS
+# ============================================================
+
+def get_spock_status():
+
+    conn = psycopg2.connect(**DB_CONFIG)
+
+    cursor = conn.cursor(
+        cursor_factory=RealDictCursor
+    )
+
+    cursor.execute("""
+        SELECT
+            subscription_name,
+            status,
+            provider_node
+        FROM spock.sub_show_status();
+    """)
+
+    rows = cursor.fetchall()
+
+    cursor.close()
+    conn.close()
+
+    return rows
+
+
+# ============================================================
+# WEBSOCKET JWT VERIFICATION
+# ============================================================
+
+def verify_websocket_token(token: str):
+
+    try:
+
+        payload = jwt.decode(
+            token,
+            SECRET_KEY,
+            algorithms=[ALGORITHM],
+        )
+
+        return payload
+
+    except ExpiredSignatureError:
+        return None
+
+    except InvalidTokenError:
+        return None
+
+
+# ============================================================
+# WEBSOCKET
+#
+# ws://localhost:8000/ws/spock-status?token=JWT_TOKEN
+# ============================================================
+
+@app.websocket("/ws/spock-status")
+async def websocket_spock_status(
+    websocket: WebSocket,
+    token: str = Query(...),
+):
+
+    token_data = verify_websocket_token(token)
+
+    if token_data is None:
+
+        await websocket.close(
+            code=1008,
+            reason="Invalid or expired JWT",
+        )
+
+        return
+
+    await websocket.accept()
+
+    try:
+
+        while True:
+
+            try:
+
+                subscription_status = (
+                    await asyncio.to_thread(
+                        get_spock_status
+                    )
+                )
+
+                await websocket.send_json({
+                    "database_reachable": True,
+                    "subscriptions": subscription_status,
+                })
+
+            except Exception as e:
+
+                await websocket.send_json({
+                    "database_reachable": False,
+                    "subscriptions": [],
+                    "error": str(e),
+                })
+
+            await asyncio.sleep(2)
+
+    except WebSocketDisconnect:
+
+        print(
+            f"WebSocket client disconnected: "
+            f"{token_data.get('sub')}"
+        )
+
+
+# ============================================================
+# GET SUBSCRIPTIONS
+# JWT PROTECTED
+# ============================================================
+
+@app.get("/spock/subscriptions")
+def get_spock_subscription_status(
+    token_data: dict = Depends(verify_token),
+):
+
+    conn = None
+    cursor = None
+
+    try:
+
+        conn = psycopg2.connect(**DB_CONFIG)
+
+        cursor = conn.cursor(
+            cursor_factory=RealDictCursor
+        )
+
+        cursor.execute("""
+            SELECT *
+            FROM spock.sub_show_status();
+        """)
+
+        rows = cursor.fetchall()
+
+        return {
+            "success": True,
+            "count": len(rows),
+            "subscriptions": rows,
+            "authenticated_as": token_data["sub"],
+        }
+
+    except Exception as e:
+
+        raise HTTPException(
+            status_code=500,
+            detail=str(e),
+        )
+
+    finally:
+
+        if cursor:
+            cursor.close()
+
+        if conn:
+            conn.close()