feat:
--Added JWT Token session
This commit is contained in:
george committed
1 parent
df12fb5188
commit
85f2b1884b
2 files changed
+642
-9
No files matched your search
@@ -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()
|
||||
Reference in new issue
Block a user