From 603aed49be3cc649462e6fbbca264e54ec731bea Mon Sep 17 00:00:00 2001
From: Wabifocus <keyboardstaff@gmail.com>
Date: Mon, 1 Sep 2025 02:14:22 -0700
Subject: [PATCH] Refactor: Update backend logic

---
 run_ui.py | 47 ++++++++++++++++++++++++++++++-----------------
 1 file changed, 30 insertions(+), 17 deletions(-)

diff --git a/run_ui.py b/run_ui.py
index 442feee2d..d96ad70bd 100644
--- a/run_ui.py
+++ b/run_ui.py
@@ -6,8 +6,7 @@ import socket
 import struct
 from functools import wraps
 import threading
-from flask import Flask, request, Response, session
-from flask_basicauth import BasicAuth
+from flask import Flask, request, Response, session, redirect, url_for, render_template_string
 import initialize
 from python.helpers import files, git, mcp_server, fasta2a_server
 from python.helpers.files import get_abs_path
@@ -43,7 +42,7 @@ webapp.config.update(
 lock = threading.Lock()
 
 # Set up basic authentication for UI and API but not MCP
-basic_auth = BasicAuth(webapp)
+# basic_auth = BasicAuth(webapp)
 
 
 def is_loopback_address(address):
@@ -79,7 +78,6 @@ def is_loopback_address(address):
                     return False
         return True
 
-
 def requires_api_key(f):
     @wraps(f)
     async def decorated(*args, **kwargs):
@@ -121,21 +119,17 @@ def requires_auth(f):
     @wraps(f)
     async def decorated(*args, **kwargs):
         user = dotenv.get_dotenv_value("AUTH_LOGIN")
-        password = dotenv.get_dotenv_value("AUTH_PASSWORD")
-        if user and password:
-            auth = request.authorization
-            if not auth or not (auth.username == user and auth.password == password):
-                return Response(
-                    "Could not verify your access level for that URL.\n"
-                    "You have to login with proper credentials",
-                    401,
-                    {"WWW-Authenticate": 'Basic realm="Login Required"'},
-                )
+        # If no auth is configured, just proceed
+        if not user:
+            return await f(*args, **kwargs)
+        
+        if not session.get('authenticated'):
+            return redirect(url_for('login'))
+        
         return await f(*args, **kwargs)
 
     return decorated
 
-
 def csrf_protect(f):
     @wraps(f)
     async def decorated(*args, **kwargs):
@@ -149,6 +143,26 @@ def csrf_protect(f):
 
     return decorated
 
+@webapp.route("/login", methods=["GET", "POST"])
+async def login():
+    error = None
+    if request.method == 'POST':
+        user = dotenv.get_dotenv_value("AUTH_LOGIN")
+        password = dotenv.get_dotenv_value("AUTH_PASSWORD")
+        
+        if request.form['username'] == user and request.form['password'] == password:
+            session['authenticated'] = True
+            return redirect(url_for('serve_index'))
+        else:
+            error = 'Invalid Credentials. Please try again.'
+            
+    login_page_content = files.read_file("webui/login.html")
+    return render_template_string(login_page_content, error=error)
+
+@webapp.route("/logout")
+async def logout():
+    session.pop('authenticated', None)
+    return redirect(url_for('login'))
 
 # handle default address, load index
 @webapp.route("/", methods=["GET"])
@@ -170,7 +184,6 @@ async def serve_index():
     )
     return index
 
-
 def run():
     PrintStyle().print("Initializing framework...")
 
@@ -267,4 +280,4 @@ def init_a0():
 if __name__ == "__main__":
     runtime.initialize()
     dotenv.load_dotenv()
-    run()
+    run()
\ No newline at end of file