Przeglądaj źródła

250311-1代码。

baoyubo 3 miesięcy temu
rodzic
commit
e9aa92a6f9
4 zmienionych plików z 69 dodań i 34 usunięć
  1. 27 2
      common/security.py
  2. 5 0
      exceptions.py
  3. 6 1
      main.py
  4. 31 31
      routers/api/__init__.py

+ 27 - 2
common/security.py

@@ -4,8 +4,10 @@ from datetime import datetime, timedelta
 import jwt
 from passlib.context import CryptContext
 from sqlalchemy.orm import Session
-from models.base import AppInfo
-from exceptions import TokenException
+from database import get_db,get_db_local
+import traceback
+from models import *
+from exceptions import TokenException,RoleException
 from config import settings
 from extensions import logger
 
@@ -29,6 +31,29 @@ def valid_access_token(Authorization: str = Header(..., alias="Authorization"))
     return int(user_id)
 
 
+def valid_access_token_role(Authorization: str = Header(..., alias="Authorization")) -> int:
+    try:
+        access_token = Authorization.removeprefix("Bearer ")
+
+        token_exception = TokenException()
+        payload = jwt.decode(access_token, settings.SECRET_KEY, algorithms=[settings.ALGORITHM])
+        print(payload, payload.get("sub"))
+        user_id: str = payload.get("sub")
+        logger.info('sub user_id: {}', user_id)
+        role_list = ["superadmin","super_ld","super_worker"]
+        db= get_db_local()
+        role_id_list = [info.role_id for info in db.query(SysRole).filter(SysRole.role_key.in_(role_list)).all()]
+        if db.query(SysUserRole).filter(SysUserRole.role_id.in_(role_id_list),SysUserRole.user_id==user_id).first() is None:
+            raise RoleException(errcode=4003, errmsg="权限不够")
+    except RoleException:
+        raise
+    except Exception:
+        # 处理异常
+        traceback.print_exc()
+        raise token_exception
+
+    return int(user_id)
+
 def valid_websocket_token(Authorization: str ) -> int:  #= Header(..., alias="sec-websocket-protocol")
     # 目前小屏测试还不能用登录功能,暂时先这样 2024/11/03
     # def valid_access_token(Authorization: str = Header("Bearer eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxIiwiZXhwIjoyMDM5Njk2ODMzfQ.Rhd38oo_S1odjg0xnT4n31cCWCAAPXGb8y_V2XcgqzQ"))->int:

+ 5 - 0
exceptions.py

@@ -39,4 +39,9 @@ class YzyException(Exception):
     def __init__(self, errcode: int, errmsg: str):
         super().__init__()
         self.errcode = errcode
+        self.errmsg = errmsg
+class RoleException(Exception):
+    def __init__(self, errcode: int, errmsg: str):
+        super().__init__(f"Error {errcode}: {errmsg}")
+        self.errcode = errcode
         self.errmsg = errmsg

+ 6 - 1
main.py

@@ -67,7 +67,12 @@ async def exception_handler(request: Request, exc: TokenException):
         status_code=200,
         content={"errcode": "4001", "errmsg":"验证不通过"}
     )
-
+@app.exception_handler(RoleException)
+async def role_exception_handler(request: Request, exc: RoleException):
+    return JSONResponse(
+        status_code=403,  # 使用 403 Forbidden 表示权限不足
+        content={"errcode": exc.errcode, "errmsg": exc.errmsg}
+    )
 @app.exception_handler(OperationalError)
 async def exception_handler(request: Request, exc: OperationalError):
     return JSONResponse(

+ 31 - 31
routers/api/__init__.py

@@ -3,7 +3,7 @@
 from fastapi import APIRouter, Request, Depends
 from starlette.templating import Jinja2Templates
 from fastapi.responses import RedirectResponse
-from common.security import valid_access_token
+from common.security import valid_access_token,valid_access_token_role
 from .login import router as login_router
 from . import gateway
 from . import dataAnalysis
@@ -52,43 +52,43 @@ router.include_router(auth.router, prefix="/auth")
 router.include_router(qrcode.router, prefix="/qrcode")
 router.include_router(login_router)
 
-router.include_router(system.router, prefix="/system", dependencies=[Depends(valid_access_token)])
+router.include_router(system.router, prefix="/system", dependencies=[Depends(valid_access_token_role)])
 
-router.include_router(gateway.router, prefix="/gateway", dependencies=[Depends(valid_access_token)])
-router.include_router(dataAnalysis.router, prefix="/dataAnalysis", dependencies=[Depends(valid_access_token)])
-router.include_router(resourceMonitoring.router, prefix="/resource", dependencies=[Depends(valid_access_token)])
-router.include_router(jobs.router, prefix="/jobs", dependencies=[Depends(valid_access_token)])
-router.include_router(OneShareProxyHandler.router, prefix="/oneShare", dependencies=[Depends(valid_access_token)])
-router.include_router(upload_file.router, prefix="/file", dependencies=[Depends(valid_access_token)])
-router.include_router(videoResource.router, prefix="/videoResource", dependencies=[Depends(valid_access_token)])
-router.include_router(Knowledge.router, prefix="/knowledge", dependencies=[Depends(valid_access_token)])
-router.include_router(taskRegistration.router, prefix="/taskRegistration", dependencies=[Depends(valid_access_token)])
-router.include_router(emergencyPlans.router, prefix="/emergency_plan", dependencies=[Depends(valid_access_token)])
-router.include_router(riskMonitor.router, prefix="/risk_monitor", dependencies=[Depends(valid_access_token)])
-router.include_router(temperaturePrecipitation.router, prefix="/temperature_precipitation", dependencies=[Depends(valid_access_token)])
-router.include_router(layerConfiguration.router, prefix="/layerConfiguration", tags=["图层配置"], dependencies=[Depends(valid_access_token)])
-router.include_router(dataFilling.router, prefix="/dataFilling", dependencies=[Depends(valid_access_token)])
-router.include_router(resourceProvison.router, prefix="/resource_provison", dependencies=[Depends(valid_access_token)])
-router.include_router(hazardStandards.router, prefix="/hazardStandards", dependencies=[Depends(valid_access_token)])
+router.include_router(gateway.router, prefix="/gateway", dependencies=[Depends(valid_access_token_role)])
+router.include_router(dataAnalysis.router, prefix="/dataAnalysis", dependencies=[Depends(valid_access_token_role)])
+router.include_router(resourceMonitoring.router, prefix="/resource", dependencies=[Depends(valid_access_token_role)])
+router.include_router(jobs.router, prefix="/jobs", dependencies=[Depends(valid_access_token_role)])
+router.include_router(OneShareProxyHandler.router, prefix="/oneShare", dependencies=[Depends(valid_access_token_role)])
+router.include_router(upload_file.router, prefix="/file", dependencies=[Depends(valid_access_token_role)])
+router.include_router(videoResource.router, prefix="/videoResource", dependencies=[Depends(valid_access_token_role)])
+router.include_router(Knowledge.router, prefix="/knowledge", dependencies=[Depends(valid_access_token_role)])
+router.include_router(taskRegistration.router, prefix="/taskRegistration", dependencies=[Depends(valid_access_token_role)])
+router.include_router(emergencyPlans.router, prefix="/emergency_plan", dependencies=[Depends(valid_access_token_role)])
+router.include_router(riskMonitor.router, prefix="/risk_monitor", dependencies=[Depends(valid_access_token_role)])
+router.include_router(temperaturePrecipitation.router, prefix="/temperature_precipitation", dependencies=[Depends(valid_access_token_role)])
+router.include_router(layerConfiguration.router, prefix="/layerConfiguration", tags=["图层配置"], dependencies=[Depends(valid_access_token_role)])
+router.include_router(dataFilling.router, prefix="/dataFilling", dependencies=[Depends(valid_access_token_role)])
+router.include_router(resourceProvison.router, prefix="/resource_provison", dependencies=[Depends(valid_access_token_role)])
+router.include_router(hazardStandards.router, prefix="/hazardStandards", dependencies=[Depends(valid_access_token_role)])
 
-router.include_router(eventManagement.router, prefix="/event_management", tags=["事件管理"], dependencies=[Depends(valid_access_token)])
-router.include_router(spatialAnalysis.router, prefix="/spatial_analysis", tags=["空间分析"], dependencies=[Depends(valid_access_token)])
-router.include_router(pattern.router, prefix="/pattern", tags=["实时测绘"], dependencies=[Depends(valid_access_token)])
+router.include_router(eventManagement.router, prefix="/event_management", tags=["事件管理"], dependencies=[Depends(valid_access_token_role)])
+router.include_router(spatialAnalysis.router, prefix="/spatial_analysis", tags=["空间分析"], dependencies=[Depends(valid_access_token_role)])
+router.include_router(pattern.router, prefix="/pattern", tags=["实时测绘"], dependencies=[Depends(valid_access_token_role)])
 
-router.include_router(rainfall.router, prefix="/rainfall", tags=["雨情分析"], dependencies=[Depends(valid_access_token)])
+router.include_router(rainfall.router, prefix="/rainfall", tags=["雨情分析"], dependencies=[Depends(valid_access_token_role)])
 
-router.include_router(infoPublish.router, prefix="/info_publish", dependencies=[Depends(valid_access_token)])
+router.include_router(infoPublish.router, prefix="/info_publish", dependencies=[Depends(valid_access_token_role)])
 
-router.include_router(riskManagement.router, prefix="/riskManagement", tags=["风险防控"], dependencies=[Depends(valid_access_token)])
+router.include_router(riskManagement.router, prefix="/riskManagement", tags=["风险防控"], dependencies=[Depends(valid_access_token_role)])
 
-router.include_router(onlineRollCall.router, prefix="/online_roll_call", tags=["在线点名"], dependencies=[Depends(valid_access_token)])
+router.include_router(onlineRollCall.router, prefix="/online_roll_call", tags=["在线点名"], dependencies=[Depends(valid_access_token_role)])
 
-router.include_router(dutyManagement.router, prefix="/duty_management", tags=["值班管理"], dependencies=[Depends(valid_access_token)])
+router.include_router(dutyManagement.router, prefix="/duty_management", tags=["值班管理"], dependencies=[Depends(valid_access_token_role)])
 
-router.include_router(companyManagement.router, prefix="/companyManagement", tags=["企业管理"], dependencies=[Depends(valid_access_token)]) #企业画像
+router.include_router(companyManagement.router, prefix="/companyManagement", tags=["企业管理"], dependencies=[Depends(valid_access_token_role)]) #企业画像
 
-router.include_router(comprehensive_search.router, prefix="/comprehensive/search", tags=["全局搜索"], dependencies=[Depends(valid_access_token)])
-router.include_router(ThreeProofingResponsible.router, prefix="/ThreeProofingResponsible", tags=["三防责任人管理"], dependencies=[Depends(valid_access_token)])
-router.include_router(city.router, prefix="/city", tags=["区划"], dependencies=[Depends(valid_access_token)])
+router.include_router(comprehensive_search.router, prefix="/comprehensive/search", tags=["全局搜索"], dependencies=[Depends(valid_access_token_role)])
+router.include_router(ThreeProofingResponsible.router, prefix="/ThreeProofingResponsible", tags=["三防责任人管理"], dependencies=[Depends(valid_access_token_role)])
+router.include_router(city.router, prefix="/city", tags=["区划"], dependencies=[Depends(valid_access_token_role)])
 
-router.include_router(yst.router, prefix="/yst", tags=["粤商通"], dependencies=[Depends(valid_access_token)])
+router.include_router(yst.router, prefix="/yst", tags=["粤商通"], dependencies=[Depends(valid_access_token_role)])