Create encrypt_image.py

encrypt_image.py

@@ -0,0 +1,206 @@
+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}')