feat :
--FAST API AGENT
This commit is contained in:
commit
ded28ab88e
5 files changed
+206
No files matched your search
@@ -0,0 +1,128 @@
|
||||
import os
|
||||
import uvicorn
|
||||
from dotenv import load_dotenv
|
||||
from fastapi import FastAPI,Depends, Header, HTTPException, status
|
||||
import psycopg2
|
||||
from jwt import ExpiredSignatureError, InvalidTokenError
|
||||
import jwt
|
||||
load_dotenv()
|
||||
app = FastAPI()
|
||||
|
||||
SECRET_KEY = os.getenv("SECRET_KEY")
|
||||
ALGORITHM = os.getenv("ALGORITHM", "HS256")
|
||||
|
||||
|
||||
|
||||
|
||||
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",
|
||||
)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@app.get("/")
|
||||
def root():
|
||||
return {"message": "PostgreSQL Agent is running"}
|
||||
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
host = os.getenv("HOST", "127.0.0.1")
|
||||
port = int(os.getenv("PORT", 9000))
|
||||
|
||||
uvicorn.run(
|
||||
"main:app",
|
||||
host=host,
|
||||
port=port,
|
||||
reload=True
|
||||
)
|
||||
|
||||
|
||||
|
||||
DB_CONFIG = {
|
||||
"host": "172.16.10.178",
|
||||
"port": 5444,
|
||||
"dbname": "postgres",
|
||||
"user": "postgres",
|
||||
"password": "postgres",
|
||||
"connect_timeout": 3,
|
||||
}
|
||||
|
||||
NODE_NAME = "node3"
|
||||
|
||||
|
||||
@app.get("/node/status")
|
||||
def node_status(token_data=Depends(verify_token)):
|
||||
|
||||
try:
|
||||
conn = psycopg2.connect(**DB_CONFIG)
|
||||
|
||||
conn.close()
|
||||
|
||||
return {
|
||||
"node": NODE_NAME,
|
||||
"online": True,
|
||||
"status": "online"
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
|
||||
return {
|
||||
"node": NODE_NAME,
|
||||
"online": False,
|
||||
"status": "offline",
|
||||
"error": str(e)
|
||||
}
|
||||
Reference in new issue
Block a user