diff --git a/tests/tools/test_green_screen_chromakey.py b/tests/tools/test_green_screen_chromakey.py index ab4bcf3a7f..6d1743cfc9 100644 --- a/tests/tools/test_green_screen_chromakey.py +++ b/tests/tools/test_green_screen_chromakey.py @@ -114,3 +114,56 @@ def _pixel(p, x, y): corner = _pixel(out, 10, 10) # was green -> keyed to dark background assert center[0] > 150 and center[1] < 80, f"subject lost, center={center!r}" assert corner[0] < 60 and corner[1] < 60, f"green not keyed, corner={corner!r}" + + +@pytest.mark.parametrize( + "invalid_color", + [ + "#000000,drawbox=x=0:y=0:w=10:h=10:color=red", + "0x00FF00:size=1x1", + "red", + "black", + "00112", + "#1234567", + "invalid_hex", + "; rm -rf /", + "#00FF00[fg];", + ], +) +def test_validate_hex_color_rejects_malformed_and_injections(invalid_color): + tool = GreenScreenProcessor() + with pytest.raises(ValueError, match="Invalid bg_color"): + tool._validate_hex_color(invalid_color) + + +@pytest.mark.parametrize( + ("valid_color", "expected_hex"), + [ + ("#0E172A", "0E172A"), + ("0E172A", "0E172A"), + ("#00ff00", "00FF00"), + ("#fff", "FFFFFF"), + ("000", "000000"), + ("#0E172AFF", "0E172AFF"), + ], +) +def test_validate_hex_color_normalizes_valid_hex(valid_color, expected_hex): + tool = GreenScreenProcessor() + assert tool._validate_hex_color(valid_color) == expected_hex + + +def test_execute_rejects_invalid_bg_color(tmp_path): + input_file = tmp_path / "input.mp4" + input_file.write_bytes(b"dummy") + output_file = tmp_path / "output.mp4" + + tool = GreenScreenProcessor() + result = tool.execute( + { + "input_path": str(input_file), + "output_path": str(output_file), + "bg_color": "#000000,drawbox=color=red", + } + ) + assert not result.success + assert "Invalid bg_color" in result.error diff --git a/tools/video/green_screen_processor.py b/tools/video/green_screen_processor.py index fdc9b7bf8d..352734f202 100644 --- a/tools/video/green_screen_processor.py +++ b/tools/video/green_screen_processor.py @@ -15,6 +15,7 @@ import json import os import platform +import re import shutil import time from pathlib import Path @@ -121,7 +122,12 @@ def execute(self, inputs: dict[str, Any]) -> ToolResult: method = inputs.get("method", "auto") fps = inputs.get("fps", 15) - bg_color = inputs.get("bg_color", "#0E172A") + bg_color_raw = inputs.get("bg_color", "#0E172A") + try: + bg_hex = self._validate_hex_color(bg_color_raw) + bg_color = f"#{bg_hex}" + except ValueError as e: + return ToolResult(success=False, error=str(e)) max_frames = inputs.get("max_frames", 0) start = time.time() @@ -452,7 +458,7 @@ def _process_chromakey( Applies chromakey to remove green, then composites onto bg_color. """ - bg_hex = bg_color.lstrip("#") + bg_hex = self._validate_hex_color(bg_color) # Convert hex to FFmpeg color format ffmpeg_bg = f"0x{bg_hex}" @@ -542,7 +548,7 @@ def _process_rembg( return False # Parse bg_color hex to RGB - bg_hex = bg_color.lstrip("#") + bg_hex = self._validate_hex_color(bg_color) bg_r = int(bg_hex[0:2], 16) bg_g = int(bg_hex[2:4], 16) bg_b = int(bg_hex[4:6], 16) @@ -605,6 +611,34 @@ def _reconstruct_video( ] self.run_command(cmd, timeout=600) + @staticmethod + def _validate_hex_color(bg_color: str) -> str: + """Validate and normalize a background hex color string. + + Accepts 3, 6, or 8 character hex strings with optional leading "#". + Returns a canonical normalized uppercase hex string (e.g., "0E172A"). + Raises ValueError if the input is not a valid hex color or contains + any injection characters. + """ + if not isinstance(bg_color, str): + raise ValueError(f"Invalid bg_color: expected string, got {type(bg_color).__name__}") + + s = bg_color.strip().removeprefix("#") + + if not re.fullmatch(r"[0-9a-fA-F]+", s): + raise ValueError( + f"Invalid bg_color: '{bg_color}'. Expected hex color string (e.g., '#0E172A')." + ) + + if len(s) == 3: + return "".join(c * 2 for c in s).upper() + elif len(s) in (6, 8): + return s.upper() + else: + raise ValueError( + f"Invalid bg_color: '{bg_color}'. Expected 3, 6, or 8 hex digits." + ) + @staticmethod def _cleanup_dir(dir_path: Path) -> None: """Recursively remove a temp directory."""