adicao de seguranca para criar admin

UnBTV · Sep 4, 2024 · 0e6700d · 0e6700d
1 parent ea6f7be
commit 0e6700d
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/src/controller/authController.py b/src/controller/authController.py
@@ -8,6 +8,7 @@
 from datetime import datetime, timedelta
 from constants import errorMessages
 from starlette.responses import JSONResponse
+from utils import security, enumeration
 
 from domain import userSchema, authSchema
 from repository import userRepository
@@ -130,7 +131,7 @@ async def validate_account(data: authSchema.AccountValidation, db: Session = Dep
 
  # cadastro da senha de admin / role do admin
 @auth.post('/admin-setup')
-async def admin_setup(data: authSchema.AdminSetup, db: Session = Depends(get_db)):
+async def admin_setup(data: authSchema.AdminSetup, db: Session = Depends(get_db), token: dict = Depends(security.verify_token_admin)):
     user = userRepository.get_user_by_email(db, data.email)
     if not user:
       raise HTTPException(status_code=404, detail=errorMessages.USER_NOT_FOUND)
@@ -146,7 +147,7 @@ async def admin_setup(data: authSchema.AdminSetup, db: Session = Depends(get_db)
     return JSONResponse(status_code=200, content={"status": "success"})
 
 @auth.post('/super-admin-setup')
-async def super_admin_setup(data: authSchema.AdminSetup, db: Session = Depends(get_db)):
+async def super_admin_setup(data: authSchema.AdminSetup, db: Session = Depends(get_db), token: dict = Depends(security.verify_token_admin)):
     user = userRepository.get_user_by_email(db, data.email)
     if not user:
       raise HTTPException(status_code=404, detail=errorMessages.USER_NOT_FOUND)

diff --git a/src/utils/security.py b/src/utils/security.py
@@ -37,9 +37,20 @@ def create_access_token(data: dict):
 def verify_token(token: str = Depends(oauth2_scheme)):
   try:
     payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+    # print(payload["role"])
     return payload
   except JWTError:
     raise HTTPException(status_code=401, detail=errorMessages.INVALID_TOKEN)
+
+def verify_token_admin(token: str = Depends(oauth2_scheme)):
+  try:
+    payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+    if payload["role"] == "ADMIN":
+      return payload
+    else:
+      raise HTTPException(status_code=401, detail=errorMessages.INVALID_TOKEN)
+  except JWTError:
+    raise HTTPException(status_code=401, detail=errorMessages.INVALID_TOKEN)
 
 def generate_six_digit_number_code():
   return secrets.randbelow(900000) + 100000