Merge pull request #4915 from QualiaRain/fix/cli-api-guards

Fix crashes on malformed inputs in cli tools and API base64 helper

cli/api-caption.py

@@ -20,7 +20,10 @@ exclude = ['a', 'in', 'on', 'out', 'at', 'the', 'and', 'with', 'next', 'to', 'it
 
 def decode(encoding):
     if encoding.startswith("data:image/"):
-        encoding = encoding.split(";")[1].split(",")[1]
+        parts = encoding.split(";", 1)
+        if len(parts) == 2:
+            parts2 = parts[1].split(",", 1)
+            encoding = parts2[1] if len(parts2) == 2 else parts2[0]
     return Image.open(io.BytesIO(base64.b64decode(encoding)))
 
 

cli/image-palette.py

@@ -21,8 +21,8 @@ grid = importlib.import_module('image-grid').grid
 
 def color_to_df(param):
     colors_pre_list = str(param).replace('([(','').split(', (')[0:-1]
-    df_rgb = [i.split('), ')[0] + ')' for i in colors_pre_list]
-    df_percent = [i.split('), ')[1].replace(')','') for i in colors_pre_list]
+    df_rgb = [i.split('), ')[0] + ')' for i in colors_pre_list if len(i.split('), ')) >= 2]
+    df_percent = [i.split('), ')[1].replace(')','') for i in colors_pre_list if len(i.split('), ')) >= 2]
     #convert RGB to HEX code
     df_color_up = [rgb2hex(int(i.split(", ")[0].replace("(","")),
                            int(i.split(", ")[1]),

cli/process.py

@@ -21,7 +21,7 @@ all_images_by_type = {}
 
 
 class Result():
-    def __init__(self, typ: str, fn: str, tag: str | None = None, requested: list = []):
+    def __init__(self, typ: str, fn: str, tag: str | None = None, requested: list | None = None):
         self.type = typ
         self.input = fn
         self.output = ''
@@ -32,7 +32,7 @@ class Result():
         self.tag = tag
         self.tags = []
         self.ops = []
-        self.steps = requested
+        self.steps = requested if requested is not None else []
 
 
 def detect_blur(image: Image.Image):
@@ -163,7 +163,9 @@ def caption_image(res: Result, tag: str | None = None):
                 for t in res.tag.split(',')[::-1]:
                     tags.insert(0, t.strip())
     pos = 0 if len(tags) == 0 else 1
-    tags.insert(pos, caption.split(' ')[1])
+    words = caption.split(' ')
+    if len(words) > 1:
+        tags.insert(pos, words[1])
     tags = [t for t in tags if len(t) > 2]
     if len(tags) > options.process.tag_limit:
         tags = tags[:options.process.tag_limit]

cli/video-extract.py

@@ -15,7 +15,10 @@ def probe(src: str):
     cmd = f"ffprobe -hide_banner -loglevel 0 -print_format json -show_format -show_streams \"{src}\""
     result = subprocess.run(cmd, shell = True, capture_output = True, text = True, check = True)
     data = json.loads(result.stdout)
-    stream = [x for x in data['streams'] if x["codec_type"] == "video"][0]
+    video_streams = [x for x in data['streams'] if x["codec_type"] == "video"]
+    if not video_streams:
+        return None
+    stream = video_streams[0]
     fmt = data['format'] if 'format' in data else {}
     res = {**stream, **fmt}
     video = Map({

modules/api/helpers.py

@@ -34,7 +34,10 @@ def decode_base64_to_image(encoding, quiet=False):
     if isinstance(encoding, str) and encoding.startswith("upload:"):
         return _resolve_upload_ref(encoding, quiet)
     if encoding.startswith("data:image/"):
-        encoding = encoding.split(";")[1].split(",")[1]
+        parts = encoding.split(";", 1)
+        if len(parts) == 2:
+            parts2 = parts[1].split(",", 1)
+            encoding = parts2[1] if len(parts2) == 2 else parts2[0]
     try:
         decoded = base64.b64decode(encoding)
         data = io.BytesIO(decoded)