146 lines
2.9 KiB
Python
146 lines
2.9 KiB
Python
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
|
|
from pydantic import BaseModel
|
|
|
|
|
|
load_dotenv()
|
|
app = FastAPI()
|
|
|
|
SECRET_KEY = os.getenv("SECRET_KEY")
|
|
ALGORITHM = os.getenv("ALGORITHM", "HS256")
|
|
|
|
|
|
class NodeStatusRequest(BaseModel):
|
|
node_name: str
|
|
host: str
|
|
port: int
|
|
dbname: str
|
|
user: str
|
|
password: str
|
|
|
|
|
|
|
|
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
|
|
)
|
|
|
|
|
|
|
|
#{
|
|
# "host": "172.16.10.178",
|
|
# "node_name = "node3",
|
|
# "port": 5444,
|
|
# "dbname": "postgres",
|
|
# "user": "postgres",
|
|
# "password": "postgres",
|
|
# "connect_timeout": 3,
|
|
# }
|
|
|
|
|
|
@app.get("/node/status")
|
|
def node_status(data: NodeStatusRequest,token_data=Depends(verify_token)):
|
|
|
|
try:
|
|
db_config = {
|
|
"host": data.host,
|
|
"port": data.port,
|
|
"dbname": data.dbname,
|
|
"user": data.user,
|
|
"password": data.password,
|
|
"connect_timeout": 3,
|
|
}
|
|
conn = psycopg2.connect(**db_config)
|
|
|
|
conn.close()
|
|
|
|
return {
|
|
"node": data.node_name,
|
|
"online": True,
|
|
"status": "online"
|
|
}
|
|
|
|
except Exception as e:
|
|
|
|
return {
|
|
"node": data.node_name,
|
|
"online": False,
|
|
"status": "offline",
|
|
"error": str(e)
|
|
} |