diff --git a/main.py b/main.py index ff0b7d0..b862ba1 100644 --- a/main.py +++ b/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(): ) + + diff --git a/main_experimental.py b/main_experimental.py new file mode 100644 index 0000000..4f03ceb --- /dev/null +++ b/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() \ No newline at end of file