Underground-Digital/Workflow-Engine

1from functools import wraps2 3from flask import request4from flask_restful import Resource5from werkzeug.exceptions import BadRequest, NotFound, Unauthorized6 7from controllers.web.error import WebSSOAuthRequiredError8from extensions.ext_database import db9from libs.passport import PassportService10from models.model import App, EndUser, Site11from services.enterprise.enterprise_service import EnterpriseService12from services.feature_service import FeatureService13 14 15def validate_jwt_token(view=None):16    def decorator(view):17        @wraps(view)18        def decorated(*args, **kwargs):19            app_model, end_user = decode_jwt_token()20 21            return view(app_model, end_user, *args, **kwargs)22 23        return decorated24 25    if view:26        return decorator(view)27    return decorator28 29 30def decode_jwt_token():31    system_features = FeatureService.get_system_features()32    app_code = request.headers.get("X-App-Code")33    try:34        auth_header = request.headers.get("Authorization")35        if auth_header is None:36            raise Unauthorized("Authorization header is missing.")37 38        if " " not in auth_header:39            raise Unauthorized("Invalid Authorization header format. Expected 'Bearer <api-key>' format.")40 41        auth_scheme, tk = auth_header.split(None, 1)42        auth_scheme = auth_scheme.lower()43 44        if auth_scheme != "bearer":45            raise Unauthorized("Invalid Authorization header format. Expected 'Bearer <api-key>' format.")46        decoded = PassportService().verify(tk)47        app_code = decoded.get("app_code")48        app_model = db.session.query(App).filter(App.id == decoded["app_id"]).first()49        site = db.session.query(Site).filter(Site.code == app_code).first()50        if not app_model:51            raise NotFound()52        if not app_code or not site:53            raise BadRequest("Site URL is no longer valid.")54        if app_model.enable_site is False:55            raise BadRequest("Site is disabled.")56        end_user = db.session.query(EndUser).filter(EndUser.id == decoded["end_user_id"]).first()57        if not end_user:58            raise NotFound()59 60        _validate_web_sso_token(decoded, system_features, app_code)61 62        return app_model, end_user63    except Unauthorized as e:64        if system_features.sso_enforced_for_web:65            app_web_sso_enabled = EnterpriseService.get_app_web_sso_enabled(app_code).get("enabled", False)66            if app_web_sso_enabled:67                raise WebSSOAuthRequiredError()68 69        raise Unauthorized(e.description)70 71 72def _validate_web_sso_token(decoded, system_features, app_code):73    app_web_sso_enabled = False74 75    # Check if SSO is enforced for web, and if the token source is not SSO, raise an error and redirect to SSO login76    if system_features.sso_enforced_for_web:77        app_web_sso_enabled = EnterpriseService.get_app_web_sso_enabled(app_code).get("enabled", False)78        if app_web_sso_enabled:79            source = decoded.get("token_source")80            if not source or source != "sso":81                raise WebSSOAuthRequiredError()82 83    # Check if SSO is not enforced for web, and if the token source is SSO,84    # raise an error and redirect to normal passport login85    if not system_features.sso_enforced_for_web or not app_web_sso_enabled:86        source = decoded.get("token_source")87        if source and source == "sso":88            raise Unauthorized("sso token expired.")89 90 91class WebApiResource(Resource):92    method_decorators = [validate_jwt_token]93