diff --git a/README.md b/README.md index 761dd6d..40248a1 100644 --- a/README.md +++ b/README.md @@ -41,6 +41,12 @@ | 🎬 MuAPI Generate (Generic) | 🎬 MuAPI | Any endpoint + raw JSON payload | | 🎬 MuAPI Save Video | 🎬 MuAPI | Download video → disk + IMAGE frames | +### Save Video safety + +The **MuAPI Save Video** node accepts only HTTPS output URLs from `cdn.muapi.ai`; redirects are +rejected. Its `save_subfolder` is confined to ComfyUI's configured output directory, and +`filename_prefix` must be a single filename component. + ## Installation **Via ComfyUI Manager:** diff --git a/muapi_video_saver_node.py b/muapi_video_saver_node.py index 57c70d2..c838bd3 100644 --- a/muapi_video_saver_node.py +++ b/muapi_video_saver_node.py @@ -1,10 +1,12 @@ """MuAPI Video Saver — downloads a video URL, saves to disk, returns frames.""" import os +from pathlib import Path +from urllib.parse import urlparse + import numpy as np import requests import torch -from PIL import Image try: import folder_paths @@ -15,6 +17,64 @@ except ImportError: return os.path.join(os.path.expanduser("~"), "comfyui_output") +ALLOWED_VIDEO_HOSTS = frozenset({"cdn.muapi.ai"}) + + +def _validate_video_url(video_url): + """Return a safe MuAPI CDN URL or raise for untrusted input.""" + if not isinstance(video_url, str): + raise ValueError("Invalid video URL") + + url = video_url.strip() + try: + parsed = urlparse(url) + hostname = parsed.hostname + port = parsed.port + except ValueError as exc: + raise ValueError("Invalid video URL") from exc + + if ( + parsed.scheme != "https" + or hostname is None + or hostname.lower() not in ALLOWED_VIDEO_HOSTS + or port not in (None, 443) + or parsed.username is not None + or parsed.password is not None + ): + raise ValueError("Video URL must be an HTTPS MuAPI CDN URL") + return url + + +def _confined_output_dir(save_subfolder): + """Resolve a user-selected subfolder without leaving ComfyUI's output directory.""" + if not isinstance(save_subfolder, str): + raise ValueError("Save subfolder must be a relative path") + + output_root = Path(folder_paths.get_output_directory()).resolve(strict=False) + requested_dir = (output_root / save_subfolder).resolve(strict=False) + try: + requested_dir.relative_to(output_root) + except ValueError as exc: + raise ValueError("Save subfolder must remain inside the output directory") from exc + return output_root, requested_dir + + +def _validate_filename_prefix(filename_prefix): + """Validate the filename component used for saved videos.""" + if not isinstance(filename_prefix, str): + raise ValueError("Filename prefix must be a single filename component") + + prefix = filename_prefix.strip() + if ( + not prefix + or prefix in {".", ".."} + or any(separator in prefix for separator in ("/", "\\", ":")) + or any(ord(character) < 32 for character in prefix) + ): + raise ValueError("Filename prefix must be a single filename component") + return prefix + + class MuAPIVideoSaver: @classmethod def INPUT_TYPES(cls): @@ -35,25 +95,40 @@ class MuAPIVideoSaver: def run(self, video_url, save_subfolder, filename_prefix, frame_load_cap=0, skip_first_frames=0, select_every_nth=1): - if not video_url or not video_url.strip().startswith("http"): - return self._err("Invalid URL") - out_dir = os.path.join(folder_paths.get_output_directory(), save_subfolder) - os.makedirs(out_dir, exist_ok=True) - n = 1 - fp = os.path.join(out_dir, f"{filename_prefix}_{n:05d}.mp4") - while os.path.exists(fp): - n += 1 - fp = os.path.join(out_dir, f"{filename_prefix}_{n:05d}.mp4") try: + safe_url = _validate_video_url(video_url) + output_root, out_dir = _confined_output_dir(save_subfolder) + safe_prefix = _validate_filename_prefix(filename_prefix) + os.makedirs(out_dir, exist_ok=True) + n = 1 + fp = os.path.join(out_dir, f"{safe_prefix}_{n:05d}.mp4") + while os.path.exists(fp): + n += 1 + fp = os.path.join(out_dir, f"{safe_prefix}_{n:05d}.mp4") + print(f"[MuAPI VideoSaver] Downloading {video_url[:80]}...") - r = requests.get(video_url, stream=True, timeout=300) + r = requests.get( + safe_url, + stream=True, + timeout=300, + allow_redirects=False, + ) + if 300 <= r.status_code < 400: + raise ValueError("MuAPI CDN redirects are not allowed") r.raise_for_status() with open(fp, "wb") as fh: for chunk in r.iter_content(8192): - if chunk: fh.write(chunk) + if chunk: + fh.write(chunk) frames, count = self._load(fp, frame_load_cap, skip_first_frames, select_every_nth) fname = os.path.basename(fp) - preview = {"filename": fname, "subfolder": save_subfolder, "type": "output", "format": "video/mp4"} + relative_subfolder = os.path.relpath(out_dir, output_root) + preview = { + "filename": fname, + "subfolder": "" if relative_subfolder == "." else relative_subfolder, + "type": "output", + "format": "video/mp4", + } print(f"[MuAPI VideoSaver] Saved {fname} — {count} frames") return {"ui": {"gifs": [preview]}, "result": (frames, fp, count)} except Exception as e: @@ -66,15 +141,21 @@ class MuAPIVideoSaver: vc = cv2.VideoCapture(path) while True: ret, frame = vc.read() - if not ret: break - if raw < skip: raw += 1; continue + if not ret: + break + if raw < skip: + raw += 1 + continue if (raw - skip) % nth == 0: rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 - frames.append(rgb); loaded += 1 - if cap > 0 and loaded >= cap: break + frames.append(rgb) + loaded += 1 + if cap > 0 and loaded >= cap: + break raw += 1 vc.release() - if not frames: raise RuntimeError("No frames") + if not frames: + raise RuntimeError("No frames") return torch.from_numpy(np.stack(frames)), len(frames) except Exception as e: print(f"[MuAPI VideoSaver] frame load error: {e}") diff --git a/test_muapi_video_saver_security.py b/test_muapi_video_saver_security.py new file mode 100644 index 0000000..77a6186 --- /dev/null +++ b/test_muapi_video_saver_security.py @@ -0,0 +1,135 @@ +import sys +import tempfile +import unittest +from pathlib import Path +from types import ModuleType +from unittest.mock import patch + +# Keep the security tests runnable in a lightweight checkout without ComfyUI's +# tensor runtime. The production node still imports the real torch package. +if "torch" not in sys.modules: + try: + import torch # noqa: F401 + except ModuleNotFoundError: + torch_stub = ModuleType("torch") + torch_stub.zeros = lambda *args, **kwargs: None + sys.modules["torch"] = torch_stub + +import muapi_video_saver_node as saver + + +class FakeResponse: + def __init__(self, status_code=200, headers=None, chunks=None): + self.status_code = status_code + self.headers = headers or {} + self._chunks = chunks or [b"video"] + + def raise_for_status(self): + if self.status_code >= 400: + raise RuntimeError(f"HTTP {self.status_code}") + + def iter_content(self, chunk_size): + return iter(self._chunks) + + +class MuAPIVideoSaverSecurityTests(unittest.TestCase): + def test_output_subfolder_cannot_escape_output_root(self): + with tempfile.TemporaryDirectory() as temp_dir: + with patch.object(saver.folder_paths, "get_output_directory", return_value=temp_dir): + _, safe_dir = saver._confined_output_dir("nested/muapi") + self.assertEqual(safe_dir, Path(temp_dir).resolve() / "nested/muapi") + + with self.assertRaisesRegex(ValueError, "inside the output directory"): + saver._confined_output_dir("../../outside") + + def test_filename_prefix_is_one_path_component(self): + self.assertEqual(saver._validate_filename_prefix(" muapi "), "muapi") + for unsafe_prefix in ("../escape", "nested/name", r"nested\\name", "C:escape", ".."): + with self.subTest(unsafe_prefix=unsafe_prefix): + with self.assertRaises(ValueError): + saver._validate_filename_prefix(unsafe_prefix) + + def test_video_url_requires_the_muapi_cdn(self): + self.assertEqual( + saver._validate_video_url(" https://cdn.muapi.ai/outputs/video.mp4 "), + "https://cdn.muapi.ai/outputs/video.mp4", + ) + for unsafe_url in ( + "http://cdn.muapi.ai/outputs/video.mp4", + "https://example.com/video.mp4", + "https://cdn.muapi.ai.evil.example/video.mp4", + "https://cdn.muapi.ai@evil.example/video.mp4", + ): + with self.subTest(unsafe_url=unsafe_url): + with self.assertRaises(ValueError): + saver._validate_video_url(unsafe_url) + + def test_run_rejects_untrusted_url_before_network_request(self): + with patch.object(saver.requests, "get") as get: + result = saver.MuAPIVideoSaver().run( + "https://example.com/video.mp4", + "muapi_videos", + "muapi", + ) + + get.assert_not_called() + self.assertIn("HTTPS MuAPI CDN", result["ui"]["text"][0]) + + def test_run_rejects_escape_subfolder_before_network_request(self): + with tempfile.TemporaryDirectory() as temp_dir: + with patch.object(saver.folder_paths, "get_output_directory", return_value=temp_dir): + with patch.object(saver.requests, "get") as get: + result = saver.MuAPIVideoSaver().run( + "https://cdn.muapi.ai/outputs/video.mp4", + "../../outside", + "muapi", + ) + + get.assert_not_called() + self.assertIn("inside the output directory", result["ui"]["text"][0]) + + def test_run_disables_redirects_for_allowlisted_url(self): + with tempfile.TemporaryDirectory() as temp_dir: + response = FakeResponse(status_code=302, headers={"Location": "https://evil.example/video.mp4"}) + with patch.object(saver.folder_paths, "get_output_directory", return_value=temp_dir): + with patch.object(saver.requests, "get", return_value=response) as get: + result = saver.MuAPIVideoSaver().run( + "https://cdn.muapi.ai/outputs/video.mp4", + "muapi_videos", + "muapi", + ) + + get.assert_called_once_with( + "https://cdn.muapi.ai/outputs/video.mp4", + stream=True, + timeout=300, + allow_redirects=False, + ) + self.assertIn("redirects are not allowed", result["ui"]["text"][0]) + + def test_run_saves_allowlisted_video_inside_output_root(self): + with tempfile.TemporaryDirectory() as temp_dir: + response = FakeResponse() + with patch.object(saver.folder_paths, "get_output_directory", return_value=temp_dir): + with patch.object(saver.requests, "get", return_value=response) as get: + with patch.object(saver.MuAPIVideoSaver, "_load", return_value=("frames", 1)): + result = saver.MuAPIVideoSaver().run( + "https://cdn.muapi.ai/outputs/video.mp4", + "muapi_videos", + "muapi", + ) + + saved_path = Path(result["result"][1]) + self.assertEqual(saved_path.parent, Path(temp_dir).resolve() / "muapi_videos") + self.assertEqual(result["ui"]["gifs"][0]["subfolder"], "muapi_videos") + self.assertTrue(saved_path.is_file()) + get.assert_called_once_with( + "https://cdn.muapi.ai/outputs/video.mp4", + stream=True, + timeout=300, + allow_redirects=False, + ) + + +if __name__ == "__main__": + unittest.main()