dify

dify/api/controllers/console/init_validate.py

@@ -0,0 +1,80 @@
+import os
+
+from flask import session
+from flask_restx import Resource, fields, reqparse
+from sqlalchemy import select
+from sqlalchemy.orm import Session
+
+from configs import dify_config
+from extensions.ext_database import db
+from libs.helper import StrLen
+from models.model import DifySetup
+from services.account_service import TenantService
+
+from . import console_ns
+from .error import AlreadySetupError, InitValidateFailedError
+from .wraps import only_edition_self_hosted
+
+
+@console_ns.route("/init")
+class InitValidateAPI(Resource):
+    @console_ns.doc("get_init_status")
+    @console_ns.doc(description="Get initialization validation status")
+    @console_ns.response(
+        200,
+        "Success",
+        model=console_ns.model(
+            "InitStatusResponse",
+            {"status": fields.String(description="Initialization status", enum=["finished", "not_started"])},
+        ),
+    )
+    def get(self):
+        """Get initialization validation status"""
+        init_status = get_init_validate_status()
+        if init_status:
+            return {"status": "finished"}
+        return {"status": "not_started"}
+
+    @console_ns.doc("validate_init_password")
+    @console_ns.doc(description="Validate initialization password for self-hosted edition")
+    @console_ns.expect(
+        console_ns.model(
+            "InitValidateRequest",
+            {"password": fields.String(required=True, description="Initialization password", max_length=30)},
+        )
+    )
+    @console_ns.response(
+        201,
+        "Success",
+        model=console_ns.model("InitValidateResponse", {"result": fields.String(description="Operation result")}),
+    )
+    @console_ns.response(400, "Already setup or validation failed")
+    @only_edition_self_hosted
+    def post(self):
+        """Validate initialization password"""
+        # is tenant created
+        tenant_count = TenantService.get_tenant_count()
+        if tenant_count > 0:
+            raise AlreadySetupError()
+
+        parser = reqparse.RequestParser().add_argument("password", type=StrLen(30), required=True, location="json")
+        input_password = parser.parse_args()["password"]
+
+        if input_password != os.environ.get("INIT_PASSWORD"):
+            session["is_init_validated"] = False
+            raise InitValidateFailedError()
+
+        session["is_init_validated"] = True
+        return {"result": "success"}, 201
+
+
+def get_init_validate_status():
+    if dify_config.EDITION == "SELF_HOSTED":
+        if os.environ.get("INIT_PASSWORD"):
+            if session.get("is_init_validated"):
+                return True
+
+            with Session(db.engine) as db_session:
+                return db_session.execute(select(DifySetup)).scalar_one_or_none()
+
+    return True