273 lines
9.1 KiB
Python
273 lines
9.1 KiB
Python
import asyncio
|
|
import os
|
|
import uuid
|
|
from contextlib import asynccontextmanager
|
|
from hashlib import sha256
|
|
from typing import List, Optional
|
|
|
|
import httpx
|
|
import uvicorn
|
|
from fastapi import Depends, FastAPI, HTTPException, Request
|
|
from fastapi.responses import JSONResponse
|
|
from fastapi.security.api_key import APIKeyHeader
|
|
from pydantic import BaseModel
|
|
from starlette.exceptions import HTTPException as StarletteHTTPException
|
|
|
|
from db_module import close_connection_pool, create_connection_pool, get_db, start_healthcheck_thread
|
|
from metrics_server import REQUEST_COUNTER
|
|
from secret_manager import cleanup_secret_files, get_secret
|
|
from secret_manager.fernet_encryption import encrypt_token
|
|
from simple_logger_handler import LOG_LEVEL, setup_logger
|
|
from validator import is_valid_platform, is_valid_token, verify_api_key
|
|
|
|
logger = setup_logger(__name__)
|
|
api_key_header = APIKeyHeader(name="X-API-Key")
|
|
|
|
class PullResponse(BaseModel):
|
|
messages: list[dict]
|
|
count: int
|
|
has_more: bool
|
|
|
|
def hash_token(token: str) -> str:
|
|
logger.debug(f"Hashing token: {token}")
|
|
result = sha256(token.encode()).hexdigest()
|
|
logger.debug(f"Hashed token result: {result}")
|
|
return result
|
|
|
|
class TokenRequest(BaseModel):
|
|
token : str
|
|
platform : str
|
|
app_ver : str
|
|
locale : Optional[str] = None
|
|
topics : Optional[List[str]] = None
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
logger.info("Starting application...")
|
|
|
|
backend_api_internal_url = os.getenv(
|
|
"BACKEND_API_INTERNAL_URL",
|
|
"http://backend-api-internal.app-notifications.svc.cluster.local:8101",
|
|
)
|
|
api_key_backend_api = get_secret(os.getenv("BACKEND_API_KEY",""),"/etc/secrets/api_key")
|
|
|
|
try:
|
|
logger.info("[FastAPI] Creating MySQL connection pool...")
|
|
create_connection_pool("ApiPool")
|
|
|
|
start_healthcheck_thread()
|
|
logger.info("[FastAPI] MySQL healthcheck thread started.")
|
|
|
|
app.state.http_client = httpx.AsyncClient (
|
|
base_url=backend_api_internal_url,
|
|
headers={"X-API-Key-Internal": api_key_backend_api},
|
|
timeout=httpx.Timeout(15.0,connect=3.0)
|
|
)
|
|
logger.info("[FastAPI] HTTP client initialized.")
|
|
|
|
SECRET_PATHS = frozenset({
|
|
"/etc/secrets/api_key",
|
|
"/etc/secrets/encryption_key",
|
|
"/etc/secrets/db_username",
|
|
"/etc/secrets/db_password"
|
|
})
|
|
cleanup_secret_files(SECRET_PATHS)
|
|
|
|
yield
|
|
|
|
finally:
|
|
logger.info("[FastAPI] Closing MySQL connection pool...")
|
|
close_connection_pool()
|
|
|
|
await app.state.http_client.aclose()
|
|
logger.info("[FastAPI] Closing Http client")
|
|
|
|
|
|
api = FastAPI(
|
|
title="Device Token Management",
|
|
description="API for requesting tokens",
|
|
version="1.0.0",
|
|
lifespan=lifespan
|
|
)
|
|
|
|
|
|
@api.middleware("http")
|
|
async def prometheus_middleware(request: Request, call_next):
|
|
status = 500
|
|
logger.debug(f"[FastAPI] Incoming request: {request.method} {request.url.path}")
|
|
try:
|
|
response = await call_next(request)
|
|
status = response.status_code
|
|
logger.debug(f"[FastAPI] Request processed with status {status}")
|
|
except Exception as e:
|
|
logger.exception(f"[FastAPI] Exception during request processing: {e}")
|
|
raise
|
|
finally:
|
|
REQUEST_COUNTER.labels(request.method, request.url.path, status).inc()
|
|
logger.debug(f"[FastAPI] Prometheus counter incremented for {request.method} {request.url.path} {status}")
|
|
return response
|
|
|
|
@api.get("/health", tags=["Health"])
|
|
def return_health(db=Depends(get_db)):
|
|
logger.debug("[FastAPI] Health check initiated")
|
|
try:
|
|
cursor = db.cursor()
|
|
cursor.execute("SELECT 1")
|
|
cursor.fetchone()
|
|
db_status = "ok"
|
|
logger.debug("[FastAPI] Database health check passed")
|
|
except Exception as e:
|
|
logger.error(f"[FastAPI] Health check DB failed: {e}")
|
|
db_status = "error"
|
|
|
|
return JSONResponse(
|
|
status_code=200 if db_status == "ok" else 500,
|
|
content={"status": db_status, "message": "Service is running"}
|
|
)
|
|
|
|
|
|
def verify_api_key_dependency(db=Depends(get_db), api_key: str = Depends(api_key_header)) -> int:
|
|
cursor = db.cursor()
|
|
cursor.execute("SELECT user_id, api_key FROM users WHERE status = 'active'")
|
|
for user_id, hashed_key in cursor.fetchall():
|
|
if verify_api_key(api_key=api_key, hashed=hashed_key):
|
|
logger.debug(f"[FastAPI] API key verified for user_id={user_id}")
|
|
return user_id
|
|
logger.warning("[FastAPI] API key verification failed")
|
|
raise HTTPException(status_code=403, detail="Unauthorized")
|
|
|
|
@api.exception_handler(StarletteHTTPException)
|
|
async def custom_http_exception_handler(request,exc):
|
|
logger.debug(f"[FastAPI] Handling HTTP exception: {exc}")
|
|
if exc.status_code == 404:
|
|
return JSONResponse(
|
|
status_code=401,
|
|
content={"detail": "Unauthorized"}
|
|
)
|
|
return JSONResponse(
|
|
status_code=exc.status_code,
|
|
content={"detail": exc.detail}
|
|
)
|
|
|
|
@api.post("/register-token")
|
|
def register_token(
|
|
request_data: TokenRequest,
|
|
db = Depends(get_db),
|
|
user_id: int = Depends(verify_api_key_dependency)
|
|
):
|
|
logger.info(f"[FastAPI] Registering token for user_id={user_id}, platform={request_data.platform}")
|
|
logger.debug(f"[FastAPI] Request data: {request_data}")
|
|
if not is_valid_platform(request_data.platform) or not is_valid_token(request_data.token):
|
|
logger.warning(f"[FastAPI] Invalid platform or token for user_id={user_id}")
|
|
raise HTTPException(status_code=403,detail="Unathorized")
|
|
|
|
secure_token = encrypt_token(request_data.token)
|
|
hashed_token = hash_token(request_data.token)
|
|
|
|
try:
|
|
cursor = db.cursor()
|
|
cursor.execute(
|
|
"SELECT * FROM device_tokens WHERE user_id=%s AND hashed_token=%s",
|
|
(user_id,hashed_token))
|
|
existing = cursor.fetchone()
|
|
logger.debug(f"[FastAPI] Existing token fetched: {existing}")
|
|
|
|
if existing:
|
|
logger.debug(f"[FastAPI] Updating existing token for user_id={user_id}")
|
|
cursor.execute("""
|
|
UPDATE device_tokens
|
|
SET platform=%s, app_ver=%s, locale=%s, topics=%s, last_seen_at=NOW()
|
|
WHERE user_id=%s AND hashed_token=%s
|
|
""",(request_data.platform,
|
|
request_data.app_ver,
|
|
request_data.locale,
|
|
request_data.topics,
|
|
user_id,
|
|
hashed_token
|
|
))
|
|
else:
|
|
token_id = str(uuid.uuid4())
|
|
logger.debug(f"[FastAPI] Creating new entry user_id={user_id}, token_id={token_id}")
|
|
cursor.execute("""
|
|
INSERT INTO device_tokens
|
|
(token_id, user_id, platform, token, hashed_token, status, app_ver, locale, topics, created_at)
|
|
VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,NOW())
|
|
""", (token_id,
|
|
user_id,
|
|
request_data.platform,
|
|
secure_token,
|
|
hashed_token,
|
|
'active',
|
|
request_data.app_ver,
|
|
request_data.locale,
|
|
request_data.topics
|
|
))
|
|
|
|
db.commit()
|
|
logger.info(f"[FastAPI] Success: Registering token for user_id={user_id}, platform={request_data.platform}")
|
|
except Exception as e:
|
|
logger.exception(f"[FastAPI] Failed to register token for user_id={user_id}: {e}")
|
|
raise HTTPException(status_code=500, detail=str(e)) from e
|
|
|
|
return {"status":"registered"}
|
|
|
|
|
|
@api.post("/unregister-token")
|
|
def unregister_token(
|
|
request_data: TokenRequest,
|
|
db = Depends(get_db),
|
|
user_id: int = Depends(verify_api_key_dependency)
|
|
):
|
|
logger.info(f"[FastAPI] Unregistering token for user_id={user_id}, platform={request_data.platform}")
|
|
hashed_token = hash_token(request_data.token)
|
|
logger.debug(f"[FastAPI] Hashed token for removal: {hashed_token}")
|
|
try:
|
|
cursor = db.cursor()
|
|
cursor.execute("""
|
|
UPDATE device_tokens
|
|
SET status=%s, last_seen_at=NOW()
|
|
WHERE hashed_token=%s
|
|
""", ('expired', hashed_token))
|
|
except Exception as e:
|
|
logger.exception(f"[FastAPI] Failed to unregister token for user_id={user_id}: {e}")
|
|
raise HTTPException(status_code=500, detail=str(e)) from e
|
|
logger.info(f"[FastAPI] Success: Unregistering token for user_id={user_id}, platform={request_data.platform}")
|
|
return {"status":"unregistered"}
|
|
|
|
@api.get("/internal/pull-notifications", response_model=PullResponse)
|
|
async def pull_notifications(
|
|
request: Request,
|
|
user_id: int = Depends(verify_api_key_dependency)
|
|
):
|
|
logger.info(f"[FastAPI] User id {user_id} is requesting messages")
|
|
client: httpx.AsyncClient = request.app.state.http_client
|
|
|
|
try:
|
|
resp = await client.post("/internal/pull-notifications",json={"user_id": user_id})
|
|
resp.raise_for_status()
|
|
data = PullResponse.model_validate(resp.json())
|
|
except httpx.HTTPStatusError as e:
|
|
logger.error(f"[FastAPI] Upstream returned {e.response.status_code} for user_id={user_id}")
|
|
raise HTTPException(status_code=502, detail="Backend-API-Internal unavailable") from e
|
|
except ValueError as e:
|
|
logger.error(f"[FastAPI] Backend-API-Internal call failed with ValueError for user_id={user_id}: {e!r}")
|
|
raise HTTPException(status_code=502, detail="Backend-API-Internal unavailable") from e
|
|
|
|
logger.info(f"[FastAPI] Returning {len(data.messages)} message(s) for user_id={user_id}, has_more={data.has_more}")
|
|
return data
|
|
|
|
async def start_servers():
|
|
logger.info("[Server] Starting main API (port 8100) and metrics server (port 9000)...")
|
|
config_main = uvicorn.Config("main:api", host="0.0.0.0", port=8100, log_level=LOG_LEVEL.lower(), log_config=None)
|
|
config_metrics = uvicorn.Config("metrics_server:metrics_api", host="0.0.0.0", port=9000, log_level=LOG_LEVEL.lower(),log_config=None)
|
|
|
|
server_main = uvicorn.Server(config_main)
|
|
server_metrics = uvicorn.Server(config_metrics)
|
|
|
|
await asyncio.gather(server_main.serve(), server_metrics.serve())
|
|
logger.info("[Server] Both servers have stopped.")
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(start_servers())
|