Added timestamp logging

Changed database lookup to look for a hashed token (that stays the same) and not a encrypted token by the vault which changes

hvac_handler.py

@@ -25,11 +25,3 @@ def decrypt_token(ciphertext: str) -> str:
 	return base64.b64decode(plaintext_b64).decode()
 
 
-if __name__ == "__main__":
-	#token = "fcm_or_apns_token_here"
-	token = "honk"
-	encrypted = encrypt_token(token)
-	print("Encrypted:", encrypted)
-
-	decrypted = decrypt_token(encrypted)
-	print("Decrypted:", decrypted)

logger_handler.py

@@ -0,0 +1,13 @@
+import logging
+
+def setup_logger(name: str) -> logging.Logger:
+	logger = logging.getLogger(name)
+	if not logger.handlers:
+		handler = logging.StreamHandler()
+		formatter = logging.Formatter(
+			'%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+		)
+		handler.setFormatter(formatter)
+		logger.addHandler(handler)
+		logger.setLevel(logging.DEBUG)
+	return logger

logging.yaml

@@ -0,0 +1,20 @@
+version: 1
+disable_existing_loggers: false
+formatters:
+  default:
+    format: "%(asctime)s - %(levelname)s - %(name)s - %(message)s"
+handlers:
+  default:
+    class: logging.StreamHandler
+    formatter: default
+    stream: ext://sys.stdout
+loggers:
+  uvicorn:
+    handlers: [default]
+    level: INFO
+  uvicorn.error:
+    handlers: [default]
+    level: INFO
+  uvicorn.access:
+    handlers: [default]
+    level: INFO

main.py

@@ -1,20 +1,29 @@
 from fastapi import FastAPI, Query, Depends, HTTPException, Header
+from fastapi.responses import JSONResponse
 from fastapi.security.api_key import APIKeyHeader
+from starlette.exceptions import HTTPException as StarletteHTTPException
 from typing import Optional,List,Dict
 from pydantic import BaseModel
 from validator import is_valid_platform,is_valid_token,verify_api_key
 from hvac_handler import encrypt_token
 from db import get_db
-import logging
+from logger_handler import setup_logger
 import uuid
 from rabbitmq_handler import send_message_to_rmq
+from hashlib import sha256
+import uvicorn
+from uvicorn_logging_config import LOGGING_CONFIG
 
 
-logger = logging.getLogger(__name__)
+
+logger = setup_logger(__name__)
 
 api_key_header = APIKeyHeader(name="X-API-Key")
 api_key_header_internal = APIKeyHeader(name="X-API-Key-Internal")
 
+def hash_token(token: str) -> str:
+	return sha256(token.encode()).hexdigest()
+
 class TokenRequest(BaseModel):
 	user_id : int
 	token : str
@@ -44,6 +53,17 @@ def verify_api_key_dependency(db=Depends(get_db), api_key: str = Depends(api_key
 			return user_id
 	raise HTTPException(status_code=403, detail="Unauthorized here")
 
+@api.exception_handler(StarletteHTTPException)
+async def custom_http_exception_handler(request,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(
@@ -56,26 +76,44 @@ def register_token(
 		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 token = %s", (secure_token,))
+		cursor.execute(
+			"SELECT * FROM device_tokens WHERE user_id=%s AND hashed_token=%s",
+			  (secure_token,hashed_token))
 		existing = cursor.fetchone()
 
 		if existing:
 			cursor.execute("""
 				UPDATE device_tokens
-				SET user_id=%s, platform=%s, app_ver=%s, locale=%s, topics=%s, last_seen_at=NOW()
-				WHERE token=%s
-			""", (user_id, request_data.platform, request_data.app_ver,
-				request_data.locale, request_data.topics, secure_token))
+				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.info(f"Creating new entry user_id={user_id}, token_id={token_id}")
 			cursor.execute("""
-				INSERT INTO device_tokens (token_id,user_id, platform, token, status, app_ver, locale, topics, created_at)
-				VALUES (%s,%s, %s, %s, %s, %s, %s, %s, NOW())
-			""", (token_id,user_id, request_data.platform, secure_token,'active',
-				request_data.app_ver, request_data.locale, request_data.topics))
+				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"Success: Registering token for user_id={user_id}, platform={request_data.platform}")
@@ -92,14 +130,14 @@ def unregister_token(
 	user_id: int = Depends(verify_api_key_dependency)
 ):
 	logger.info(f"Unregistering token for user_id={user_id}, platform={request_data.platform}")
-	secure_token = encrypt_token(request_data.token)
+	hashed_token = hash_token(request_data.token)
 	try:
 		cursor = db.cursor()
 		cursor.execute("""
 				UPDATE device_tokens
 				SET status=%s, last_seen_at=NOW()
-				WHERE token=%s
-			""", ('expired', secure_token))
+				WHERE hashed_token=%s
+			""", ('expired', hashed_token))
 		
 	except Exception as e:
 		raise HTTPException(status_code=500, detail=str(e))
@@ -114,4 +152,14 @@ def receive_notifications(
 	is_allowed: bool = Depends(verify_api_key_dependency_internal)
 ):
 	send_message_to_rmq(notification_data.user_id,notification_data.message)
-	return {"status": "queued"}
+	return {"status": "queued"}
+
+
+if __name__ == "__main__":
+	uvicorn.run(
+		"main:api",
+		host="0.0.0.0",
+		port=8000,
+		log_config=LOGGING_CONFIG,
+		log_level="info"
+	)

uvicorn_logging_config.py

@@ -0,0 +1,39 @@
+LOGGING_CONFIG = {
+	"version": 1,
+	"disable_existing_loggers": False,
+	"formatters": {
+		"default": {
+			"format": "%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+			"datefmt": "%Y-%m-%d %H:%M:%S",
+		}
+	},
+	"handlers": {
+		"default": {
+			"class": "logging.StreamHandler",
+			"formatter": "default",
+			"stream": "ext://sys.stdout"
+		}
+	},
+	"loggers": {
+		"": {  # root logger
+			"handlers": ["default"],
+			"level": "INFO",
+			"propagate": False
+		},
+		"uvicorn": {
+			"handlers": ["default"],
+			"level": "INFO",
+			"propagate": False
+		},
+		"uvicorn.error": {
+			"handlers": ["default"],
+			"level": "INFO",
+			"propagate": False
+		},
+		"uvicorn.access": {
+			"handlers": ["default"],
+			"level": "INFO",
+			"propagate": False
+		}
+	}
+}