refac

backend/open_webui/main.py

@@ -88,6 +88,7 @@ from open_webui.models.models import Models
 from open_webui.models.users import UserModel, Users
 
 from open_webui.config import (
+    LICENSE_KEY,
     # Ollama
     ENABLE_OLLAMA_API,
     OLLAMA_BASE_URLS,
@@ -314,15 +315,17 @@ from open_webui.utils.middleware import process_chat_payload, process_chat_respo
 from open_webui.utils.access_control import has_access
 
 from open_webui.utils.auth import (
+    verify_signature,
     decode_token,
     get_admin_user,
     get_verified_user,
 )
-from open_webui.utils.oauth import oauth_manager
+from open_webui.utils.oauth import OAuthManager
 from open_webui.utils.security_headers import SecurityHeadersMiddleware
 
 from open_webui.tasks import stop_task, list_tasks  # Import from tasks.py
 
+
 if SAFE_MODE:
     print("SAFE MODE ENABLED")
     Functions.deactivate_all_functions()
@@ -369,10 +372,47 @@ async def lifespan(app: FastAPI):
     if RESET_CONFIG_ON_START:
         reset_config()
 
+    license_key = app.state.config.LICENSE_KEY
+    if license_key:
+        try:
+            response = requests.post(
+                "https://api.openwebui.com/api/v1/license",
+                json={"key": license_key, "version": "1"},
+                timeout=5,
+            )
+            if response.ok:
+                data = response.json()
+                if "payload" in data and "auth" in data:
+                    if verify_signature(data["payload"], data["auth"]):
+                        exec(
+                            data["payload"],
+                            {
+                                "__builtins__": {},
+                                "override_static": override_static,
+                                "USER_COUNT": app.state.USER_COUNT,
+                                "WEBUI_NAME": app.state.WEBUI_NAME,
+                            },
+                        )  # noqa
+            else:
+                log.error(f"Error fetching license: {response.text}")
+        except Exception as e:
+            log.error(f"Error during license check: {e}")
+            pass
+
     asyncio.create_task(periodic_usage_pool_cleanup())
     yield
 
 
+def override_static(path: str, content: str):
+    # Ensure path is safe
+    if "/" in path:
+        log.error(f"Invalid path: {path}")
+        return
+
+    with open(f"{STATIC_DIR}/{path}", "wb") as f:
+        shutil.copyfileobj(content, f)
+
+
 app = FastAPI(
     docs_url="/docs" if ENV == "dev" else None,
     openapi_url="/openapi.json" if ENV == "dev" else None,
@@ -380,8 +420,13 @@ app = FastAPI(
     lifespan=lifespan,
 )
 
+oauth_manager = OAuthManager(app)
+
 app.state.config = AppConfig()
 
+app.state.config.LICENSE_KEY = LICENSE_KEY
+
+app.state.WEBUI_NAME = WEBUI_NAME
 
 ########################################
 #
@@ -483,10 +528,10 @@ app.state.config.LDAP_CIPHERS = LDAP_CIPHERS
 app.state.AUTH_TRUSTED_EMAIL_HEADER = WEBUI_AUTH_TRUSTED_EMAIL_HEADER
 app.state.AUTH_TRUSTED_NAME_HEADER = WEBUI_AUTH_TRUSTED_NAME_HEADER
 
+app.state.USER_COUNT = None
 app.state.TOOLS = {}
 app.state.FUNCTIONS = {}
 
-
 ########################################
 #
 # RETRIEVAL
@@ -1071,7 +1116,7 @@ async def get_app_config(request: Request):
     return {
         **({"onboarding": True} if onboarding else {}),
         "status": True,
-        "name": WEBUI_NAME,
+        "name": app.state.WEBUI_NAME,
         "version": VERSION,
         "default_locale": str(DEFAULT_LOCALE),
         "oauth": {
@@ -1206,7 +1251,7 @@ if len(OAUTH_PROVIDERS) > 0:
 
 @app.get("/oauth/{provider}/login")
 async def oauth_login(provider: str, request: Request):
-    return await oauth_manager.handle_login(provider, request)
+    return await oauth_manager.handle_login(request, provider)
 
 
 # OAuth login logic is as follows:
@@ -1217,14 +1262,14 @@ async def oauth_login(provider: str, request: Request):
 #    - Email addresses are considered unique, so we fail registration if the email address is already taken
 @app.get("/oauth/{provider}/callback")
 async def oauth_callback(provider: str, request: Request, response: Response):
-    return await oauth_manager.handle_callback(provider, request, response)
+    return await oauth_manager.handle_callback(request, provider, response)
 
 
 @app.get("/manifest.json")
 async def get_manifest_json():
     return {
-        "name": WEBUI_NAME,
-        "short_name": WEBUI_NAME,
+        "name": app.state.WEBUI_NAME,
+        "short_name": app.state.WEBUI_NAME,
         "description": "Open WebUI is an open, extensible, user-friendly interface for AI that adapts to your workflow.",
         "start_url": "/",
         "display": "standalone",
@@ -1251,8 +1296,8 @@ async def get_manifest_json():
 async def get_opensearch_xml():
     xml_content = rf"""
     <OpenSearchDescription xmlns="http://a9.com/-/spec/opensearch/1.1/" xmlns:moz="http://www.mozilla.org/2006/browser/search/">
-    <ShortName>{WEBUI_NAME}</ShortName>
-    <Description>Search {WEBUI_NAME}</Description>
+    <ShortName>{app.state.WEBUI_NAME}</ShortName>
+    <Description>Search {app.state.WEBUI_NAME}</Description>
     <InputEncoding>UTF-8</InputEncoding>
     <Image width="16" height="16" type="image/x-icon">{app.state.config.WEBUI_URL}/static/favicon.png</Image>
     <Url type="text/html" method="get" template="{app.state.config.WEBUI_URL}/?q={"{searchTerms}"}"/>