Underground-Digital/Workflow-Engine

1from datetime import datetime, timezone2 3from flask_login import current_user4from flask_restful import Resource, marshal_with, reqparse5from werkzeug.exceptions import Forbidden, NotFound6 7from constants.languages import supported_language8from controllers.console import api9from controllers.console.app.wraps import get_app_model10from controllers.console.wraps import account_initialization_required, setup_required11from extensions.ext_database import db12from fields.app_fields import app_site_fields13from libs.login import login_required14from models import Site15 16 17def parse_app_site_args():18    parser = reqparse.RequestParser()19    parser.add_argument("title", type=str, required=False, location="json")20    parser.add_argument("icon_type", type=str, required=False, location="json")21    parser.add_argument("icon", type=str, required=False, location="json")22    parser.add_argument("icon_background", type=str, required=False, location="json")23    parser.add_argument("description", type=str, required=False, location="json")24    parser.add_argument("default_language", type=supported_language, required=False, location="json")25    parser.add_argument("chat_color_theme", type=str, required=False, location="json")26    parser.add_argument("chat_color_theme_inverted", type=bool, required=False, location="json")27    parser.add_argument("customize_domain", type=str, required=False, location="json")28    parser.add_argument("copyright", type=str, required=False, location="json")29    parser.add_argument("privacy_policy", type=str, required=False, location="json")30    parser.add_argument("custom_disclaimer", type=str, required=False, location="json")31    parser.add_argument(32        "customize_token_strategy", type=str, choices=["must", "allow", "not_allow"], required=False, location="json"33    )34    parser.add_argument("prompt_public", type=bool, required=False, location="json")35    parser.add_argument("show_workflow_steps", type=bool, required=False, location="json")36    parser.add_argument("use_icon_as_answer_icon", type=bool, required=False, location="json")37    return parser.parse_args()38 39 40class AppSite(Resource):41    @setup_required42    @login_required43    @account_initialization_required44    @get_app_model45    @marshal_with(app_site_fields)46    def post(self, app_model):47        args = parse_app_site_args()48 49        # The role of the current user in the ta table must be editor, admin, or owner50        if not current_user.is_editor:51            raise Forbidden()52 53        site = db.session.query(Site).filter(Site.app_id == app_model.id).one_or_404()54 55        for attr_name in [56            "title",57            "icon_type",58            "icon",59            "icon_background",60            "description",61            "default_language",62            "chat_color_theme",63            "chat_color_theme_inverted",64            "customize_domain",65            "copyright",66            "privacy_policy",67            "custom_disclaimer",68            "customize_token_strategy",69            "prompt_public",70            "show_workflow_steps",71            "use_icon_as_answer_icon",72        ]:73            value = args.get(attr_name)74            if value is not None:75                setattr(site, attr_name, value)76 77        site.updated_by = current_user.id78        site.updated_at = datetime.now(timezone.utc).replace(tzinfo=None)79        db.session.commit()80 81        return site82 83 84class AppSiteAccessTokenReset(Resource):85    @setup_required86    @login_required87    @account_initialization_required88    @get_app_model89    @marshal_with(app_site_fields)90    def post(self, app_model):91        # The role of the current user in the ta table must be admin or owner92        if not current_user.is_admin_or_owner:93            raise Forbidden()94 95        site = db.session.query(Site).filter(Site.app_id == app_model.id).first()96 97        if not site:98            raise NotFound99 100        site.code = Site.generate_code(16)101        site.updated_by = current_user.id102        site.updated_at = datetime.now(timezone.utc).replace(tzinfo=None)103        db.session.commit()104 105        return site106 107 108api.add_resource(AppSite, "/apps/<uuid:app_id>/site")109api.add_resource(AppSiteAccessTokenReset, "/apps/<uuid:app_id>/site/access-token-reset")110