proof-of-concept
This commit is contained in:
Sida Liu
@@ -1,3 +1,5 @@
|
||||
login_new.html
|
||||
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
|
||||
@@ -1,2 +1,29 @@
|
||||
# ComfyUI-Login
|
||||
A custom node that implements basic login
|
||||
This custom node serves as a proof-of-concept to explore the implementation of basic login functionality for [ComfyUI](https://github.com/comfyanonymous/ComfyUI).
|
||||
|
||||
In the future, if proven useful, this feature might be directly integrated into either [ComfyUI](https://github.com/comfyanonymous/ComfyUI) or [ComfyUI-Manager](https://github.com/ltdrdata/ComfyUI-Manager).
|
||||
|
||||
# First thing first
|
||||
Please be aware that absolute security does not exist. This login node offers only basic protection for ComfyUI. Use it at your own risk.
|
||||
|
||||
# How to use
|
||||
|
||||
## Installation
|
||||
To install this node, you have two options:
|
||||
|
||||
1. Use [ComfyUI-Manager](https://github.com/ltdrdata/ComfyUI-Manager).
|
||||
|
||||
2. In the directory `ComfyUI/custom_nodes/`, git clone this repo.
|
||||
|
||||
## Setting Up a New Password
|
||||
For your first login, you may choose any password. This password will be encrypted and stored in a file named "PASSWORD" within the ComfyUI project folder.
|
||||
|
||||
## Resetting a Forgotten Password
|
||||
If you forget your password, you can reset it by deleting the "PASSWORD" file found in the ComfyUI project folder. After doing this, you will be able to log in again using a new password of your choice.
|
||||
|
||||
## Removing the Login Feature
|
||||
To disable the login feature, you can either:
|
||||
|
||||
1. Use the [ComfyUI-Manager](https://github.com/ltdrdata/ComfyUI-Manager) to uninstall it, or
|
||||
|
||||
2. Manually delete the ComfyUI-Login folder located in the `ComfyUI/custom_nodes/` directory. Then, restart ComfyUI for the changes to take effect.
|
||||
|
||||
+83
@@ -0,0 +1,83 @@
|
||||
import server
|
||||
import aiohttp
|
||||
from aiohttp_session import setup, get_session
|
||||
from aiohttp_session.cookie_storage import EncryptedCookieStorage
|
||||
from aiohttp import web
|
||||
import base64
|
||||
import os
|
||||
import folder_paths
|
||||
import bcrypt
|
||||
|
||||
node_dir = os.path.dirname(__file__)
|
||||
comfy_dir = os.path.dirname(folder_paths.__file__)
|
||||
password_path = os.path.join(comfy_dir, "PASSWORD")
|
||||
|
||||
# Access the PromptServer instance and its app
|
||||
prompt_server = server.PromptServer.instance
|
||||
app = prompt_server.app
|
||||
routes = prompt_server.routes
|
||||
|
||||
secret_key = base64.urlsafe_b64decode(os.getenv('SECRET_KEY', base64.urlsafe_b64encode(os.urandom(32))))
|
||||
setup(app, EncryptedCookieStorage(secret_key))
|
||||
|
||||
@routes.get("/login")
|
||||
async def get_root(request):
|
||||
session = await get_session(request)
|
||||
if ('logged_in' in session and session['logged_in']):
|
||||
raise web.HTTPFound('/')
|
||||
else:
|
||||
return web.FileResponse(os.path.join(node_dir, "login.html"))
|
||||
|
||||
# Add a route for "/login" that handles the form submission
|
||||
@routes.post("/login")
|
||||
async def login_handler(request):
|
||||
data = await request.post() # Get the data from the form
|
||||
password = data.get('password').encode('utf-8')
|
||||
if os.path.exists(password_path):
|
||||
with open(password_path, "rb") as f:
|
||||
hashed_password = f.read()
|
||||
if bcrypt.checkpw(password, hashed_password):
|
||||
session = await get_session(request)
|
||||
session['logged_in'] = True
|
||||
response = web.HTTPFound('/') # Redirect to the main page if the password is correct
|
||||
return response
|
||||
else:
|
||||
return web.HTTPFound('/login?feedback=Wrong password')
|
||||
else:
|
||||
salt = bcrypt.gensalt()
|
||||
hashed_password = bcrypt.hashpw(password, salt)
|
||||
# Write the hashed password to a file
|
||||
with open(password_path, "wb") as file:
|
||||
file.write(hashed_password)
|
||||
session = await get_session(request)
|
||||
session['logged_in'] = True
|
||||
response = web.HTTPFound('/') # Redirect to the main page if the password is correct
|
||||
return response
|
||||
|
||||
@routes.get("/logout")
|
||||
async def get_root(request):
|
||||
session = await get_session(request)
|
||||
session['logged_in'] = False
|
||||
response = web.HTTPFound('/login') # Redirect to the main page if the password is correct
|
||||
return response
|
||||
|
||||
# For loading all custom js
|
||||
WEB_DIRECTORY = "js"
|
||||
|
||||
@web.middleware
|
||||
async def check_login_status(request: web.Request, handler):
|
||||
if request.path == '/login' or request.path.endswith('.css') or request.path.endswith('.js'):
|
||||
# Skip for safe URIs
|
||||
response = await handler(request)
|
||||
return response
|
||||
session = await get_session(request)
|
||||
if ('logged_in' in session and session['logged_in']):
|
||||
response = await handler(request)
|
||||
if request.path == '/': # This avoids seeing the GUI after logging out and navigating back immediately.
|
||||
response.headers.setdefault('Cache-Control', 'no-cache')
|
||||
return response
|
||||
raise web.HTTPFound('/login')
|
||||
|
||||
app.middlewares.append(check_login_status)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
@@ -0,0 +1,20 @@
|
||||
import { app } from "../../scripts/app.js";
|
||||
|
||||
app.registerExtension({
|
||||
name: "Comfy.Password",
|
||||
init() {},
|
||||
async setup() {
|
||||
await new Promise(resolve => setTimeout(resolve, 500)); // Delay for 0.5 second before appending the Logout button to the menu, ensuring it is added last.
|
||||
|
||||
const menu = document.querySelector(".comfy-menu");
|
||||
|
||||
const logoutButton = document.createElement("button");
|
||||
logoutButton.textContent = "Logout";
|
||||
logoutButton.onclick = () => {
|
||||
localStorage.clear(); // If you use localStorage
|
||||
sessionStorage.clear(); // If you use sessionStorage
|
||||
window.location.href = "/logout";
|
||||
}
|
||||
menu.append(logoutButton);
|
||||
},
|
||||
});
|
||||
+6
-6
@@ -1,12 +1,12 @@
|
||||
# Setting Up a New Password
|
||||
When logging in for the first time, you can choose any password. This password will be encrypted and stored in a file named "PASSWORD" located within the ComfyUI project folder.
|
||||
## Setting Up a New Password
|
||||
For your first login, you may choose any password. This password will be encrypted and stored in a file named "PASSWORD" within the ComfyUI project folder.
|
||||
|
||||
# Resetting a Forgotten Password
|
||||
## Resetting a Forgotten Password
|
||||
If you forget your password, you can reset it by deleting the "PASSWORD" file found in the ComfyUI project folder. After doing this, you will be able to log in again using a new password of your choice.
|
||||
|
||||
# Removing the Login Feature
|
||||
## Removing the Login Feature
|
||||
To disable the login feature, you can either:
|
||||
|
||||
1. Use the ComfyUI-Manager to uninstall it, or
|
||||
1. Use the [ComfyUI-Manager](https://github.com/ltdrdata/ComfyUI-Manager) to uninstall it, or
|
||||
|
||||
2. Manually delete the ComfyUI-Login folder located in the ComfyUI/custom_nodes/ directory. Then, restart ComfyUI for the changes to take effect.
|
||||
2. Manually delete the ComfyUI-Login folder located in the `ComfyUI/custom_nodes/` directory. Then, restart ComfyUI for the changes to take effect.
|
||||
|
||||
+44
@@ -0,0 +1,44 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<title>ComfyUI Login</title>
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0, user-scalable=no">
|
||||
<link rel="stylesheet" type="text/css" href="./lib/litegraph.css" />
|
||||
<link rel="stylesheet" type="text/css" href="./style.css" />
|
||||
<link rel="stylesheet" type="text/css" href="./user.css" />
|
||||
</head>
|
||||
<body class="litegraph" style="text-align: center;">
|
||||
<h2>ComfyUI Login</h2>
|
||||
<div id="feedback" style="margin: 20px;"></div>
|
||||
<form action="/login" method="post">
|
||||
<label for="password">Password:</label>
|
||||
<input type="password" id="password" name="password" required>
|
||||
<br><br>
|
||||
<input type="submit" value="Submit" class="comfy-btn">
|
||||
<br><br>
|
||||
<a href="https://github.com/liusida/ComfyUI-Login/blob/main/locked_out.md" style="color:white; font-size:x-small" target="_blank">Locked Out?</a>
|
||||
</form>
|
||||
</body>
|
||||
<script>
|
||||
// Function to parse query parameters
|
||||
function getQueryParams() {
|
||||
var params = {};
|
||||
window.location.search.substring(1).split("&").forEach(function(part) {
|
||||
var item = part.split("=");
|
||||
// Ensure both key and value are defined; replace '+' with ' ' (space) in value if defined
|
||||
var key = decodeURIComponent(item[0]);
|
||||
var value = item[1] ? decodeURIComponent(item[1].replace(/\+/g, ' ')) : null;
|
||||
params[key] = value;
|
||||
});
|
||||
return params;
|
||||
}
|
||||
|
||||
var params = getQueryParams();
|
||||
var feedbackElement = document.getElementById('feedback');
|
||||
if (params.feedback) {
|
||||
feedbackElement.innerText = params.feedback;
|
||||
}
|
||||
</script>
|
||||
|
||||
</html>
|
||||
@@ -0,0 +1,2 @@
|
||||
aiohttp-session
|
||||
bcrypt
|
||||
@@ -0,0 +1,32 @@
|
||||
import bcrypt
|
||||
import getpass # Used for secure password input
|
||||
import os
|
||||
import folder_paths
|
||||
comfy_dir = os.path.dirname(folder_paths.__file__)
|
||||
password_path = os.path.join(comfy_dir, "PASSWORD")
|
||||
|
||||
def main():
|
||||
# Prompt the user for a new password securely
|
||||
password = getpass.getpass('Enter a new password: ')
|
||||
# Prompt for the password again for verification
|
||||
password_verify = getpass.getpass('Re-enter your password: ')
|
||||
|
||||
# Check if both entered passwords match
|
||||
if password == password_verify:
|
||||
# Convert the password to bytes
|
||||
password_bytes = password.encode('utf-8')
|
||||
|
||||
# Generate a salt and hash the password
|
||||
salt = bcrypt.gensalt()
|
||||
hashed_password = bcrypt.hashpw(password_bytes, salt)
|
||||
|
||||
# Write the hashed password to a file
|
||||
with open(password_path, "wb") as file:
|
||||
file.write(hashed_password)
|
||||
|
||||
print("Password has been securely stored. Now you can restart ComfyUI.")
|
||||
else:
|
||||
print("Passwords do not match. Please try again.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user