diff --git a/.env.sample b/.env.sample new file mode 100644 index 000000000..ccfd6cf46 --- /dev/null +++ b/.env.sample @@ -0,0 +1,2 @@ +SDAPI_USERNAME=admin +SDAPI_PASSWORD=test123 \ No newline at end of file diff --git a/.github/workflows/on_merge_to_master.yaml b/.github/workflows/on_merge_to_master.yaml index b221540c0..d09fcb6d5 100644 --- a/.github/workflows/on_merge_to_master.yaml +++ b/.github/workflows/on_merge_to_master.yaml @@ -21,6 +21,16 @@ jobs: with: ssh-private-key: ${{ secrets.MOONLITE_AWS_EC2_SSH_KEY }} + - name: Create .env file + working-directory: ./ + run: | + echo "SDAPI_USERNAME=${{ secrets.MOONLITE_SDAPI_USERNAME }}" >> .env + echo "SDAPI_PASSWORD=${{ secrets.MOONLITE_SDAPI_PASSWORD }}" >> .env + + - name: Copy .env file to EC2 instance + run: | + scp -o StrictHostKeyChecking=no -i ${{ secrets.MOONLITE_AWS_EC2_SSH_KEY }} .env ${{ secrets.MOONLITE_AWS_EC2_SSH_USER }}@${{ secrets.MOONLITE_AWS_EC2_SSH_HOST }}:/home/ec2-user/apps/stable-diffusion/ + - name: SSH into EC2 and deploy the app run: | ssh -o ServerAliveInterval=60 -o ServerAliveCountMax=60 -o StrictHostKeyChecking=no ${{ secrets.MOONLITE_AWS_EC2_SSH_USER }}@${{ secrets.MOONLITE_AWS_EC2_SSH_HOST }} << 'EOF' diff --git a/.gitignore b/.gitignore index 8c119e3db..d4eda3bd2 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,7 @@ __pycache__ /venv /tmp /model.ckpt +.env # /models/**/* /GFPGANv1.3.pth /gfpgan/weights/*.pth diff --git a/basic_auth_middleware.py b/basic_auth_middleware.py new file mode 100644 index 000000000..55c818b43 --- /dev/null +++ b/basic_auth_middleware.py @@ -0,0 +1,39 @@ +import base64 +from starlette.middleware.base import BaseHTTPMiddleware +from starlette.responses import Response +from fastapi import FastAPI + +class BasicAuthMiddleware(BaseHTTPMiddleware): + def __init__(self, app, username: str, password: str): + super().__init__(app) + self.username = username + self.password = password + + async def dispatch(self, request, call_next): + # Extract the Authorization header + auth_header = request.headers.get("Authorization") + if not auth_header or not auth_header.startswith("Basic "): + return self._unauthorized_response() + + try: + # Decode and split the credentials + encoded_credentials = auth_header.split(" ")[1] + # We should add a step to santize the input here + decoded_credentials = base64.b64decode(encoded_credentials).decode("utf-8") + provided_username, provided_password = decoded_credentials.split(":") + + # Check credentials + if provided_username == self.username and provided_password == self.password: + response = await call_next(request) + return response + else: + return self._unauthorized_response() + except Exception: + return self._unauthorized_response() + + def _unauthorized_response(self): + return Response( + content="Unauthorized", + status_code=401, + headers={"WWW-Authenticate": "Basic realm='api'"} + ) diff --git a/extensions/sd-webui-controlnet b/extensions/sd-webui-controlnet index 8bbbd0e55..56cec5b29 160000 --- a/extensions/sd-webui-controlnet +++ b/extensions/sd-webui-controlnet @@ -1 +1 @@ -Subproject commit 8bbbd0e55ef6e5d71b09c2de2727b36e7bc825b0 +Subproject commit 56cec5b2958edf3b1807b7e7b2b1b5186dbd2f81 diff --git a/requirements.txt b/requirements.txt index 9e2ecfe4d..44bcf09e0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -20,6 +20,7 @@ open-clip-torch piexif psutil pytorch_lightning +python-dotenv requests resize-right diff --git a/webui.py b/webui.py index 2c417168a..73d370a8b 100644 --- a/webui.py +++ b/webui.py @@ -1,5 +1,6 @@ from __future__ import annotations +from dotenv import load_dotenv import os import time @@ -9,6 +10,7 @@ from modules import initialize startup_timer = timer.startup_timer startup_timer.record("launcher") +load_dotenv() initialize.imports() @@ -30,6 +32,13 @@ def api_only(): initialize.initialize() app = FastAPI() + + # Initialize the Basic Authentication middleware + from basic_auth_middleware import BasicAuthMiddleware + USERNAME = os.getenv('SDAPI_USERNAME') + PASSWORD = os.getenv('SDAPI_PASSWORD') + app.add_middleware(BasicAuthMiddleware, username=USERNAME, password=PASSWORD) + initialize_util.setup_middleware(app) api = create_api(app)