""" Security regression tests. These tests pin behavior that, if it regresses silently, would re-introduce a previously-fixed vulnerability. They run without a real ComfyUI process by stubbing the `server.PromptServer` module. Run from the repo root: python -m pytest tests/ Or directly: python -m unittest tests.test_security """ import importlib import os import stat import sys import tempfile import types import unittest from pathlib import Path from unittest import mock REPO_ROOT = Path(__file__).resolve().parent.parent PARENT_OF_REPO = REPO_ROOT.parent # Install ghost modules for the google-auth / googleapiclient packages # BEFORE the package under test imports them. This lets the test run # without those packages installed in the environment. def _install_lib_ghosts(): """Install fake modules for runtime deps so the package can import without google-auth / googleapiclient / watchdog installed.""" for name in [ "google", "google.oauth2", "google.auth", "google.auth.transport", "googleapiclient", "google_auth_oauthlib", "watchdog", "watchdog.observers", "watchdog.events", ]: if name not in sys.modules: sys.modules[name] = types.ModuleType(name) sys.modules["google.oauth2.credentials"] = types.ModuleType("google.oauth2.credentials") sys.modules["google.oauth2.credentials"].Credentials = mock.MagicMock() sys.modules["google.oauth2.service_account"] = types.ModuleType("google.oauth2.service_account") sys.modules["google.oauth2.service_account"].Credentials = mock.MagicMock() sys.modules["google.auth.transport.requests"] = types.ModuleType("google.auth.transport.requests") sys.modules["google.auth.transport.requests"].Request = mock.MagicMock() sys.modules["googleapiclient.discovery"] = types.ModuleType("googleapiclient.discovery") sys.modules["googleapiclient.discovery"].build = mock.MagicMock() sys.modules["googleapiclient.http"] = types.ModuleType("googleapiclient.http") sys.modules["googleapiclient.http"].MediaFileUpload = mock.MagicMock() sys.modules["googleapiclient.errors"] = types.ModuleType("googleapiclient.errors") sys.modules["googleapiclient.errors"].HttpError = type("HttpError", (Exception,), {}) sys.modules["google_auth_oauthlib.flow"] = types.ModuleType("google_auth_oauthlib.flow") sys.modules["google_auth_oauthlib.flow"].Flow = mock.MagicMock() sys.modules["google_auth_oauthlib.flow"].InstalledAppFlow = mock.MagicMock() sys.modules["watchdog.observers"].Observer = mock.MagicMock() sys.modules["watchdog.events"].FileSystemEventHandler = type("FileSystemEventHandler", (object,), {}) _install_lib_ghosts() def _install_fake_promptserver(client_id, sockets=None): """Install a stubbed `server` module so the Setup Node can import it. Returns a `Mock` for `PromptServer.instance.send_sync` so tests can assert on calls. """ fake_module = types.ModuleType("server") fake_instance = mock.MagicMock() fake_instance.client_id = client_id fake_instance.sockets = sockets if sockets is not None else {} fake_instance.send_sync = mock.MagicMock() fake_module.PromptServer = mock.MagicMock() fake_module.PromptServer.instance = fake_instance sys.modules["server"] = fake_module return fake_instance.send_sync def _install_fake_gdrive_auth_manager(oauth_response=None): """Install a fake gdrive_auth_manager module so tests do not require the google-auth packages to be installed. """ fake = types.ModuleType("gdrive_auth_manager") fake.get_oauth_credentials_for_setup = mock.MagicMock( return_value=oauth_response or {"credentials": { "client_id": "stub.apps.googleusercontent.com", "client_secret": "GOCSPX-stub", "refresh_token": "1//stub", "token": "ya29.stub", "token_uri": "https://oauth2.googleapis.com/token", "scopes": ["https://www.googleapis.com/auth/drive.file"], }} ) fake.get_drive_service = mock.MagicMock() fake.get_folder_id = mock.MagicMock(return_value=None) fake.get_oauth_credentials = mock.MagicMock(return_value=None) fake.get_service_account_credentials = mock.MagicMock(return_value=None) fake.refresh_access_token = mock.MagicMock() sys.modules["gdrive_auth_manager"] = fake sys.modules["ComfyUI_DriveSendNode.gdrive_auth_manager"] = fake return fake def _import_setup_node_fresh(): """Re-import the setup node module as part of the ComfyUI_DriveSendNode package so relative imports (e.g. `from .gdrive_auth_manager import ...`) resolve correctly. """ pkg_name = "ComfyUI_DriveSendNode" mod_name = f"{pkg_name}.drivesend_setup_node" # Strip cached entries so import-time route registration sees the # current `server` mock. for cached in list(sys.modules): if cached == mod_name or cached.startswith(mod_name + "."): sys.modules.pop(cached, None) sys.modules.pop(mod_name, None) if str(PARENT_OF_REPO) not in sys.path: sys.path.insert(0, str(PARENT_OF_REPO)) return importlib.import_module(mod_name) class TestWebSocketBroadcastGuard(unittest.TestCase): """Setup must NEVER call send_sync(..., sid=None) for credential events. ComfyUI's send_sync with sid=None broadcasts to every connected WebSocket client. When PromptServer.instance.client_id is None (typical for curl/SDK submissions that omit client_id), credentials delivered via display_only would otherwise leak to any concurrent listener on /ws. """ def test_no_send_sync_when_client_id_is_none(self): send_sync = _install_fake_promptserver(client_id=None) setup_mod = _import_setup_node_fresh() with tempfile.TemporaryDirectory() as tmp: with mock.patch.object(setup_mod, "NODE_DIR", Path(tmp)): with mock.patch.dict(os.environ, {"COMFYUI_DRIVESEND_ALLOW_SETUP": "1"}, clear=False): node = setup_mod.DriveSendSetupNode() # Force a state-changing path: stash an OAuth attempt # with a usable client_id+secret+auth_code, then call # setup() with no real client_id present. setup_mod._stash_set("any-client", { "google_client_id": "stub.apps.googleusercontent.com", "google_client_secret": "GOCSPX-stub", "auth_code": "code-stub", "service_account_json": "", }) # The fake PromptServer reports client_id=None, so # _stash_consume(None) returns nothing. The OAuth # branch will still proceed but find no inputs and # return a "click Set credentials..." message # without touching send_sync. with mock.patch( "ComfyUI_DriveSendNode.gdrive_auth_manager.get_oauth_credentials_for_setup", return_value={"credentials": { "client_id": "stub.apps.googleusercontent.com", "client_secret": "GOCSPX-stub", "refresh_token": "1//stub", "token": "ya29.stub", "token_uri": "https://oauth2.googleapis.com/token", "scopes": ["https://www.googleapis.com/auth/drive.file"], }}, ): node.setup( auth_method="oauth", folder_id="stub_folder_id", storage_method="display_only", encryption_key_method="off", reconnect=False, ) # Assert: every send_sync call was scoped to a non-None sid. for call in send_sync.call_args_list: args, kwargs = call sid = args[2] if len(args) >= 3 else kwargs.get("sid") self.assertIsNotNone( sid, f"send_sync was called with sid=None: {call}. " "This would broadcast to every connected WebSocket client.", ) def test_send_sync_called_with_real_sid(self): """Sanity check: when a real client_id is present, send_sync IS called and is scoped to that sid.""" send_sync = _install_fake_promptserver( client_id="real-sid", sockets={"real-sid": object()} ) setup_mod = _import_setup_node_fresh() with tempfile.TemporaryDirectory() as tmp: with mock.patch.object(setup_mod, "NODE_DIR", Path(tmp)): with mock.patch.dict(os.environ, {"COMFYUI_DRIVESEND_ALLOW_SETUP": "1"}, clear=False): setup_mod._stash_set("real-sid", { "google_client_id": "stub.apps.googleusercontent.com", "google_client_secret": "GOCSPX-stub", "auth_code": "code-stub", "service_account_json": "", }) with mock.patch( "ComfyUI_DriveSendNode.gdrive_auth_manager.get_oauth_credentials_for_setup", return_value={"credentials": { "client_id": "stub.apps.googleusercontent.com", "client_secret": "GOCSPX-stub", "refresh_token": "1//stub", "token": "ya29.stub", "token_uri": "https://oauth2.googleapis.com/token", "scopes": ["https://www.googleapis.com/auth/drive.file"], }}, ): node = setup_mod.DriveSendSetupNode() node.setup( auth_method="oauth", folder_id="stub_folder_id", storage_method="display_only", encryption_key_method="off", reconnect=False, ) # Expect at least one send_sync delivering the credentials_ready # event scoped to real-sid. creds_calls = [ c for c in send_sync.call_args_list if c.args and c.args[0] == "drivesend_credentials_ready" ] self.assertTrue(creds_calls, "Expected drivesend_credentials_ready to be sent") for call in creds_calls: self.assertEqual(call.args[2], "real-sid") payload = call.args[1] self.assertEqual(payload.get("client_id"), "real-sid", "Payload must echo the sid for JS-side defense in depth") class TestEnvFilePermissions(unittest.TestCase): """The .env file must be created with mode 0o600. Default umask typically produces 0o644, which is world-readable on multi-tenant hosts. The Setup and AutoUploader nodes both use os.open with O_CREAT|0o600 to create the file race-free.""" def test_persist_env_writes_creates_0600(self): _install_fake_promptserver(client_id=None) setup_mod = _import_setup_node_fresh() with tempfile.TemporaryDirectory() as tmp: with mock.patch.object(setup_mod, "NODE_DIR", Path(tmp)): env_path = setup_mod._persist_env_writes( [("GOOGLE_DRIVE_FOLDER_ID", "stub")], scrub_setup_keys=False, ) self.assertTrue(env_path.exists()) mode = stat.S_IMODE(env_path.stat().st_mode) self.assertEqual( mode, 0o600, f".env file mode is {oct(mode)}, expected 0o600. " "Default-umask creation would expose secrets to other local users.", ) class TestStashHardening(unittest.TestCase): """The /drivesend/setup/stash route must enforce: client_id pattern, payload size cap, JSON content type, live-WebSocket gating. These are the hardening primitives that prevent stash poisoning.""" def test_client_id_pattern_rejects_bad_input(self): _install_fake_promptserver(client_id=None) setup_mod = _import_setup_node_fresh() good = ["abc", "abc-123", "client.id_42", "X" * 128] bad = ["", "abc 123", "abc/../etc", "X" * 129, "