Added support for authentication in --share mode (via auth.json)

Author: MoonRide303

.gitignore

@@ -41,3 +41,4 @@ notification.mp3
 /resolutions.json
 /notification.ogg
 /notification.mp3
+/auth.json

auth-example.json

@@ -0,0 +1,6 @@
+[
+    {
+        "user": "sitting-duck-1",
+        "pass": "very-bad-publicly-known-password-change-it"
+    }
+]

modules/auth.py

@@ -0,0 +1,42 @@
+import json
+import hashlib
+
+from os.path import exists
+
+
+def auth_list_to_dict(auth_list):
+    auth_dict = {}
+    for auth_data in auth_list:
+        if 'user' in auth_data:
+            if 'hash' in auth_data:
+                auth_dict |= {auth_data['user']: auth_data['hash']}
+            elif 'pass' in auth_data:
+                auth_dict |= {auth_data['user']: hashlib.sha256(bytes(auth_data['pass'], encoding='utf-8')).hexdigest()}
+    return auth_dict
+
+
+def load_auth_data(filename=None):
+    auth_dict = None
+    if filename != None and exists(filename):
+        with open(filename, encoding='utf-8') as auth_file:
+            try:
+                auth_obj = json.load(auth_file)
+                if isinstance(auth_obj, list) and len(auth_obj) > 0:
+                    auth_dict = auth_list_to_dict(auth_obj)
+            except Exception as e:
+                print('load_auth_data, e: ' + str(e))
+            finally:
+                auth_file.close()
+    return auth_dict
+
+
+auth_dict = load_auth_data('auth.json')
+
+auth_enabled = auth_dict != None
+
+
+def check_auth(user, password):
+    if user not in auth_dict:
+        return False
+    else:   
+        return hashlib.sha256(bytes(password, encoding='utf-8')).hexdigest() == auth_dict[user]

readme.md

@@ -182,6 +182,7 @@ Below things are already inside the software, and **users do not need to do anyt
 24. Support for playing audio when generation is finished (ported from SD web UI - use notification.ogg or notification.mp3).
 25. Starting generation via Ctrl-ENTER hotkey (ported from SD web UI).
 26. Support for loading models from subfolders (ported from RuinedFooocus).
+27. Support for authentication in --share mode (credentials loaded from auth.json - use auth-examle.json as a template).
 
 ## Thanks
 

update_log_mre.md

@@ -1,3 +1,7 @@
+### 2.0.18 MRE
+
+* Added support for authentication in --share mode (via auth.json).
+
 ### 2.0.14 MRE
 
 * Added support for loading models from subfolders (ported from RuinedFooocus).

webui.py

@@ -19,6 +19,7 @@
 from fastapi import FastAPI
 from modules.ui_gradio_extensions import reload_javascript
 from modules.util import get_current_log_path, get_previous_log_path
+from modules.auth import auth_enabled, check_auth
 from os.path import exists
 
 
@@ -485,4 +486,4 @@ def stop_clicked():
 
 
 app = gr.mount_gradio_app(app, shared.gradio_root, '/')
-shared.gradio_root.launch(inbrowser=True, server_name=args.listen, server_port=args.port, share=args.share)
+shared.gradio_root.launch(inbrowser=True, server_name=args.listen, server_port=args.port, share=args.share, auth=check_auth if args.share and auth_enabled else None)