ui/encrypt_image.py

import base64
import io
from pathlib import Path
from modules import shared,script_callbacks,scripts as md_scripts,images
from modules.api import api
from modules.shared import opts
from scripts.core.core import get_sha256,dencrypt_image,dencrypt_image_v2,encrypt_image_v2
from PIL import PngImagePlugin,_util,ImagePalette
from PIL import Image as PILImage
from io import BytesIO
from typing import Optional
from fastapi import FastAPI
from gradio import Blocks
from fastapi import FastAPI, Request, Response
import sys
from urllib.parse import unquote
from colorama import Fore, Back, Style

repo_dir = md_scripts.basedir()
password = getattr(shared.cmd_opts, 'encrypt_pass', None)


def hook_http_request(app: FastAPI):
    @app.middleware("http")
    async def image_dencrypt(req: Request, call_next):
        endpoint:str = req.scope.get('path', 'err')
        endpoint='/'+endpoint.strip('/')
        # 兼容无边浏览器
        if endpoint.startswith('/infinite_image_browsing/image-thumbnail') or endpoint.startswith('/infinite_image_browsing/file'):
            query_string:str = req.scope.get('query_string').decode('utf-8')
            query_string = unquote(query_string)
            if query_string and query_string.index('path=')>=0:
                query = query_string.split('&')
                path = ''
                for sub in query:
                    if sub.startswith('path='):
                        path = sub[sub.index('=')+1:]
                if path:
                    endpoint = '/file=' + path
        # 模型预览图
        if endpoint.startswith('/sd_extra_networks/thumb'):
            query_string:str = req.scope.get('query_string').decode('utf-8')
            query_string = unquote(query_string)
            if query_string and query_string.index('filename=')>=0:
                query = query_string.split('&')
                path = ''
                for sub in query:
                    if sub.startswith('filename='):
                        path = sub[sub.index('=')+1:]
                if path:
                    endpoint = '/file=' + path
        if endpoint.startswith('/file='):
            file_path = endpoint[6:] or ''
            if not file_path: return await call_next(req)
            if file_path.rfind('.') == -1: return await call_next(req)
            if not file_path[file_path.rfind('.'):]: return await call_next(req)
            if file_path[file_path.rfind('.'):].lower() in ['.png','.jpg','.jpeg','.webp','.abcd']:
                image = PILImage.open(file_path)
                pnginfo = image.info or {}
                if 'Encrypt' in pnginfo:
                    buffered = BytesIO()
                    info = PngImagePlugin.PngInfo()
                    for key in pnginfo.keys():
                        if pnginfo[key]:
                            info.add_text(key,pnginfo[key])
                    image.save(buffered, format=PngImagePlugin.PngImageFile.format, pnginfo=info)
                    decrypted_image_data = buffered.getvalue()
                    response: Response = Response(content=decrypted_image_data, media_type="image/png")
                    return response
        
        return await call_next(req)
    
def set_shared_options():
    # 传递插件状态到前端
    section = ("encrypt_image_is_enable",'图片加密' if shared.opts.localization == 'zh_CN' else "encrypt image" )
    option = shared.OptionInfo(
            default="是",
            label='是否启用了加密插件' if shared.opts.localization == 'zh_CN' else "Whether the encryption plug-in is enabled",
            section=section,
        )
    option.do_not_save = True
    shared.opts.add_option(
        "encrypt_image_is_enable",
        option,
    )
    shared.opts.data['encrypt_image_is_enable'] = "是"

def app_started_callback(_: Blocks, app: FastAPI):
    set_shared_options()
    

if PILImage.Image.__name__ != 'EncryptedImage':
    super_open = PILImage.open
    super_encode_pil_to_base64 = api.encode_pil_to_base64
    super_modules_images_save_image = images.save_image
    super_api_middleware = api.api_middleware
    class EncryptedImage(PILImage.Image):
        __name__ = "EncryptedImage"
        
        @staticmethod
        def from_image(image:PILImage.Image):
            image = image.copy()
            img = EncryptedImage()
            img.im = image.im
            img._mode = image.mode
            if image.im.mode:
                try:
                    img.mode = image.im.mode
                except Exception as e:
                    ''
            img._size = image.size
            img.format = image.format
            if image.mode in ("P", "PA"):
                if image.palette:
                    img.palette = image.palette.copy()
                else:
                    img.palette = ImagePalette.ImagePalette()
            img.info = image.info.copy()
            return img
            
        def save(self, fp, format=None, **params):
            filename = ""
            if isinstance(fp, Path):
                filename = str(fp)
            elif _util.is_path(fp):
                filename = fp
            elif fp == sys.stdout:
                try:
                    fp = sys.stdout.buffer
                except AttributeError:
                    pass
            if not filename and hasattr(fp, "name") and _util.is_path(fp.name):
                # only set the name for metadata purposes
                filename = fp.name
            
            if not filename or not password:
                # 如果没有密码或不保存到硬盘，直接保存
                super().save(fp, format = format, **params)
                return
            
            if 'Encrypt' in self.info and (self.info['Encrypt'] == 'pixel_shuffle' or self.info['Encrypt'] == 'pixel_shuffle_2'):
                super().save(fp, format = format, **params)
                return
            
            encrypt_image_v2(self, get_sha256(password))
            self.format = PngImagePlugin.PngImageFile.format
            pnginfo = params.get('pnginfo', PngImagePlugin.PngInfo())
            if not pnginfo:
                pnginfo = PngImagePlugin.PngInfo()
            pnginfo.add_text('Encrypt', 'pixel_shuffle_2')
            pnginfo.add_text('EncryptPwdSha', get_sha256(f'{get_sha256(password)}Encrypt'))
            for key in (self.info or {}).keys():
                if self.info[key]:
                    pnginfo.add_text(key,str(self.info[key]))
            params.update(pnginfo=pnginfo)
            super().save(fp, format=self.format, **params)
            # 保存到文件后解密内存内的图片，让直接在内存内使用时图片正常
            dencrypt_image_v2(self, get_sha256(password))
            


    def open(fp,*args, **kwargs):
        image = super_open(fp,*args, **kwargs)
        if password and image.format.lower() == PngImagePlugin.PngImageFile.format.lower():
            pnginfo = image.info or {}
            if 'Encrypt' in pnginfo and pnginfo["Encrypt"] == 'pixel_shuffle':
                dencrypt_image(image, get_sha256(password))
                pnginfo["Encrypt"] = None
                image = EncryptedImage.from_image(image=image)
                return image
            if 'Encrypt' in pnginfo and pnginfo["Encrypt"] == 'pixel_shuffle_2':
                dencrypt_image_v2(image, get_sha256(password))
                pnginfo["Encrypt"] = None
                image = EncryptedImage.from_image(image=image)
                return image
        return EncryptedImage.from_image(image=image)
    
    def encode_pil_to_base64(image:PILImage.Image):
        with io.BytesIO() as output_bytes:
            image.save(output_bytes, format="PNG", quality=opts.jpeg_quality)
            pnginfo = image.info or {}
            if 'Encrypt' in pnginfo and pnginfo["Encrypt"] == 'pixel_shuffle':
                dencrypt_image(image, get_sha256(password))
                pnginfo["Encrypt"] = None
            if 'Encrypt' in pnginfo and pnginfo["Encrypt"] == 'pixel_shuffle_2':
                dencrypt_image_v2(image, get_sha256(password))
                pnginfo["Encrypt"] = None
            bytes_data = output_bytes.getvalue()
        return base64.b64encode(bytes_data)
  
    def api_middleware(app: FastAPI):
        super_api_middleware(app)
        hook_http_request(app)
  
    if password:
        PILImage.Image = EncryptedImage
        PILImage.open = open
        api.encode_pil_to_base64 = encode_pil_to_base64
        api.api_middleware = api_middleware
        
if password:
    script_callbacks.on_app_started(app_started_callback)
    print(f'{Fore.GREEN}[-] Image Encryption started.{Style.RESET_ALL}')

else:
    print(f'{Fore.RED}[-] Image Encryption DISABLED.{Style.RESET_ALL}')