318 lines
10 KiB
Python
318 lines
10 KiB
Python
import os
|
|
import jwt
|
|
import uuid
|
|
from aiohttp import web
|
|
|
|
from server import PromptServer
|
|
|
|
from .utils import *
|
|
|
|
instance = PromptServer.instance
|
|
app = instance.app
|
|
routes = instance.routes
|
|
|
|
logger = Logger(LOG_FILE, LOG_LEVELS)
|
|
sanitizer = Sanitizer()
|
|
ip_filter = IPFilter(WHITELIST, BLACKLIST)
|
|
timeout = Timeout(ip_filter, BLACKLIST_AFTER_ATTEMPTS)
|
|
users_db = UsersDB(USERS_FILE)
|
|
access_control = AccessControl(users_db, instance, HISTORY_FILE)
|
|
jwt_auth = JWTAuth(
|
|
users_db, access_control, logger, SECRET_KEY, TOKEN_EXPIRE_MINUTES, TOKEN_ALGORITHM
|
|
)
|
|
access_control.sync_comfy_users()
|
|
|
|
|
|
@routes.get("/register")
|
|
async def get_register(request: web.Request) -> web.Response:
|
|
with open(os.path.join(HTML_DIR, "register.html"), "r") as f:
|
|
html_content = f.read()
|
|
|
|
if not users_db.load_users():
|
|
html_content = html_content.replace("{{ X-Admin-User }}", "true")
|
|
html_content = html_content.replace("{{ X-Admin-User }}", "false")
|
|
|
|
return web.Response(body=html_content, content_type="text/html")
|
|
|
|
|
|
@routes.post("/register")
|
|
async def post_register(request: web.Request) -> web.Response:
|
|
sanitized_data = request.get("_sanitized_data", {})
|
|
ip = get_ip(request)
|
|
new_user_username = sanitized_data.get("new_user_username")
|
|
new_user_password = sanitized_data.get("new_user_password")
|
|
username = sanitized_data.get("username")
|
|
password = sanitized_data.get("password")
|
|
|
|
username_valid, username_invalid_message = validate_username(new_user_username)
|
|
if not username_valid:
|
|
return web.json_response({"error": username_invalid_message}, status=400)
|
|
|
|
password_valid, password_invalid_message = validate_password(new_user_password)
|
|
if not password_valid:
|
|
return web.json_response({"error": password_invalid_message}, status=400)
|
|
|
|
admin_user = users_db.get_admin_user()
|
|
admin_user_id = None
|
|
|
|
if admin_user[0] and (not new_user_username or not new_user_password):
|
|
return web.json_response(
|
|
{"error": "Missing new user registration details"}, status=400
|
|
)
|
|
|
|
if admin_user[0]:
|
|
if not username or not password:
|
|
return web.json_response(
|
|
{"error": "Missing admin user authentication details"}, status=400
|
|
)
|
|
|
|
admin_user_id = admin_user[0]
|
|
|
|
if admin_user_id is not None:
|
|
if not (
|
|
users_db.get_user(username)[0] == admin_user_id
|
|
and users_db.check_username_password(username, password)
|
|
):
|
|
logger.registration_attempt(
|
|
ip, username, password, new_user_username, new_user_password
|
|
)
|
|
timeout.add_failed_attempt(ip)
|
|
return web.json_response(
|
|
{"message": "Invalid username or password"}, status=403
|
|
)
|
|
|
|
if None not in users_db.get_user(new_user_username):
|
|
return web.json_response({"error": "Username already exists"}, status=400)
|
|
|
|
new_user_id = str(uuid.uuid4())
|
|
users_db.add_user(
|
|
new_user_id,
|
|
new_user_username,
|
|
new_user_password,
|
|
not bool(admin_user_id),
|
|
)
|
|
access_control.sync_comfy_users()
|
|
|
|
logger.registration_success(
|
|
ip, new_user_username, username if admin_user_id is not None else None
|
|
)
|
|
timeout.remove_failed_attempts(ip)
|
|
return web.json_response({"message": "User registered successfully"})
|
|
|
|
|
|
@routes.get("/login")
|
|
async def get_login(request: web.Request) -> web.Response:
|
|
if not users_db.load_users():
|
|
return web.HTTPFound("/register")
|
|
|
|
token = jwt_auth.get_token_from_request(request)
|
|
if token:
|
|
return web.HTTPFound("/logout")
|
|
return web.FileResponse(os.path.join(HTML_DIR, "login.html"))
|
|
|
|
|
|
@routes.post("/login")
|
|
async def post_login(request: web.Request) -> web.Response:
|
|
sanitized_data = request.get("_sanitized_data", {})
|
|
ip = get_ip(request)
|
|
username = sanitized_data.get("username")
|
|
password = sanitized_data.get("password")
|
|
|
|
if not username or not password:
|
|
return web.json_response(
|
|
{"error": "Missing login credentials (username and password)"}, status=400
|
|
)
|
|
|
|
if users_db.check_username_password(username, password):
|
|
timeout.remove_failed_attempts(ip)
|
|
|
|
user_id, _ = users_db.get_user(username)
|
|
token = jwt_auth.create_access_token({"id": user_id, "username": username})
|
|
response = web.json_response(
|
|
{
|
|
"message": "Login successful",
|
|
"jwt_token": token,
|
|
"user_settings_id": user_id,
|
|
}
|
|
)
|
|
secure_flag = request.headers.get("X-Forwarded-Proto", "http") == "https"
|
|
response.set_cookie(
|
|
"jwt_token",
|
|
token,
|
|
max_age=TOKEN_EXPIRE_MINUTES * 60,
|
|
httponly=True,
|
|
secure=secure_flag,
|
|
samesite="Strict",
|
|
)
|
|
logger.login_success(ip, username)
|
|
return response
|
|
|
|
logger.login_attempt(ip, username, password)
|
|
timeout.add_failed_attempt(ip)
|
|
return web.json_response({"error": "Invalid username or password"}, status=401)
|
|
|
|
|
|
@routes.get("/generate_token")
|
|
async def get_generate_token(request: web.Request) -> web.Response:
|
|
if not users_db.load_users():
|
|
return web.HTTPFound("/register")
|
|
|
|
token = jwt_auth.get_token_from_request(request)
|
|
if token:
|
|
return web.HTTPFound("/logout")
|
|
return web.FileResponse(os.path.join(HTML_DIR, "generate_token.html"))
|
|
|
|
|
|
@routes.post("/generate_token")
|
|
async def post_generate_token(request: web.Request) -> web.Response:
|
|
sanitized_data = request.get("_sanitized_data", {})
|
|
ip = get_ip(request)
|
|
username = sanitized_data.get("username")
|
|
password = sanitized_data.get("password")
|
|
|
|
try:
|
|
expire_hours = int(
|
|
sanitized_data.get("expire_hours", TOKEN_EXPIRE_MINUTES / 60)
|
|
)
|
|
|
|
except ValueError:
|
|
return web.json_response(
|
|
{"error": "Expiration hours must be a number"},
|
|
status=400,
|
|
)
|
|
|
|
if expire_hours > MAX_TOKEN_EXPIRE_MINUTES / 60:
|
|
return web.json_response(
|
|
{
|
|
"error": f"Expiration hours must be smaller than {MAX_TOKEN_EXPIRE_MINUTES / 60}"
|
|
},
|
|
status=400,
|
|
)
|
|
|
|
if not username or not password:
|
|
return web.json_response(
|
|
{"error": "Missing login credentials (username and password)"}, status=400
|
|
)
|
|
|
|
if users_db.check_username_password(username, password):
|
|
timeout.remove_failed_attempts(ip)
|
|
|
|
user_id, _ = users_db.get_user(username)
|
|
token = jwt_auth.create_access_token(
|
|
{"id": user_id, "username": username}, expire_minutes=(expire_hours * 60)
|
|
)
|
|
response = web.json_response(
|
|
{
|
|
"message": "JWT Token successfully generated",
|
|
"jwt_token": token,
|
|
}
|
|
)
|
|
secure_flag = request.headers.get("X-Forwarded-Proto", "http") == "https"
|
|
response.set_cookie(
|
|
"jwt_token",
|
|
token,
|
|
max_age=expire_hours * 60 * 60,
|
|
httponly=True,
|
|
secure=secure_flag,
|
|
samesite="Strict",
|
|
)
|
|
|
|
logger.generate_success(ip, username, expire_hours)
|
|
|
|
return response
|
|
|
|
logger.generate_attempt(ip, username, password, expire_hours)
|
|
timeout.add_failed_attempt(ip)
|
|
return web.json_response({"error": "Invalid username or password"}, status=401)
|
|
|
|
|
|
@routes.get("/logout")
|
|
async def get_logout(request: web.Request) -> web.Response:
|
|
ip = get_ip(request)
|
|
free_memory = request.query.get("free_memory", "false").lower() == "true"
|
|
unload_models = request.query.get("unload_models", "false").lower() == "true"
|
|
|
|
token = jwt_auth.get_token_from_request(request)
|
|
if token and FREE_MEMORY_ON_LOGOUT:
|
|
try:
|
|
username = jwt_auth.decode_access_token(token).get("username")
|
|
if free_memory or unload_models:
|
|
if hasattr(instance, "post_free"):
|
|
json_data = {
|
|
"unload_models": unload_models,
|
|
"free_memory": free_memory,
|
|
}
|
|
mock_request = web.Request(
|
|
app=app,
|
|
method="POST",
|
|
path="/free",
|
|
headers={},
|
|
match_info={},
|
|
payload=None,
|
|
)
|
|
mock_request._post = json_data
|
|
await instance.post_free(mock_request)
|
|
logger.memory_free(ip, username, free_memory, unload_models)
|
|
|
|
logger.logout(ip, username)
|
|
except jwt.ExpiredSignatureError:
|
|
pass
|
|
except jwt.InvalidTokenError:
|
|
pass
|
|
except Exception as e:
|
|
logger.error(f"Unexpected error during logout: {e}")
|
|
|
|
response = web.HTTPFound("/login")
|
|
response.del_cookie("jwt_token", path="/")
|
|
|
|
return response
|
|
|
|
|
|
app.add_routes(
|
|
[
|
|
web.static("/account-manager/css", CSS_DIR),
|
|
web.static("/account-manager/js", JS_DIR),
|
|
web.static("/account-manager/assets", ASSETS_DIR),
|
|
]
|
|
)
|
|
|
|
if FORCE_HTTPS:
|
|
app.middlewares.append(create_https_middleware(MATCH_HEADERS))
|
|
|
|
app.middlewares.append(ip_filter.create_ip_filter_middleware())
|
|
app.middlewares.append(sanitizer.create_sanitizer_middleware())
|
|
app.middlewares.append(
|
|
timeout.create_time_out_middleware(
|
|
limited=("/login", "/register", "/generate_token")
|
|
)
|
|
)
|
|
app.middlewares.append(
|
|
jwt_auth.create_jwt_middleware(
|
|
public=("/login", "/logout", "/register", "/generate_token"),
|
|
public_prefixes=("/account-manager"),
|
|
)
|
|
)
|
|
|
|
if SEPERATE_USERS:
|
|
access_control.patch_comfy_user_manager()
|
|
access_control.patch_assets()
|
|
app.middlewares.append(access_control.create_view_access_control_middleware())
|
|
app.middlewares.append(access_control.create_folder_access_control_middleware())
|
|
|
|
access_control.patch_folder_paths()
|
|
access_control.patch_prompt_queue()
|
|
|
|
if MANAGER_ADMIN_ONLY:
|
|
app.middlewares.append(
|
|
access_control.create_manager_access_control_middleware(
|
|
manager_directory="/extensions/comfyui-manager",
|
|
manager_routes=(
|
|
"api/customnode",
|
|
"api/snapshot",
|
|
"/api/manager",
|
|
"api/comfyui_manager",
|
|
"api/externalmodel",
|
|
),
|
|
)
|
|
)
|