diff --git a/comfy_cli/command/models/models.py b/comfy_cli/command/models/models.py index ffb795dd..2f02aadd 100644 --- a/comfy_cli/command/models/models.py +++ b/comfy_cli/command/models/models.py @@ -21,6 +21,8 @@ workspace_manager = WorkspaceManager() config_manager = ConfigManager() +_CIVITAI_SUBDOMAIN_SUFFIXES = tuple(f".{h}" for h in constants.CIVITAI_ALLOWED_HOSTS) + model_path_map = { "lora": "loras", @@ -99,7 +101,7 @@ def check_civitai_url(url: str) -> tuple[bool, bool, int | None, int | None]: try: parsed = urlparse(url) host = (parsed.hostname or "").lower() - if host != "civitai.com" and not host.endswith(".civitai.com"): + if host not in constants.CIVITAI_ALLOWED_HOSTS and not host.endswith(_CIVITAI_SUBDOMAIN_SUFFIXES): return False, False, None, None p_parts = [p for p in parsed.path.split("/") if p] query = parse_qs(parsed.query) diff --git a/comfy_cli/constants.py b/comfy_cli/constants.py index 9e561bc8..9e1517ee 100644 --- a/comfy_cli/constants.py +++ b/comfy_cli/constants.py @@ -47,6 +47,7 @@ class PROC(str, Enum): CIVITAI_API_TOKEN_KEY = "civitai_api_token" CIVITAI_API_TOKEN_ENV_KEY = "CIVITAI_API_TOKEN" +CIVITAI_ALLOWED_HOSTS: tuple[str, ...] = ("civitai.com", "civitai.red") HF_API_TOKEN_KEY = "hf_api_token" HF_API_TOKEN_ENV_KEY = "HF_API_TOKEN" diff --git a/tests/comfy_cli/command/models/test_models.py b/tests/comfy_cli/command/models/test_models.py index cd70e481..13dae746 100644 --- a/tests/comfy_cli/command/models/test_models.py +++ b/tests/comfy_cli/command/models/test_models.py @@ -244,6 +244,41 @@ def test_non_evil_civitai_url(): assert check_civitai_url(url) == (False, False, None, None) +def test_valid_model_url_red_domain(): + url = "https://civitai.red/models/43331" + assert check_civitai_url(url) == (True, False, 43331, None) + + +def test_valid_model_url_red_with_query(): + url = "https://civitai.red/models/43331?modelVersionId=485088" + assert check_civitai_url(url) == (True, False, 43331, 485088) + + +def test_valid_api_download_url_red_domain(): + url = "https://civitai.red/api/download/models/1617665?type=Model&format=SafeTensor" + assert check_civitai_url(url) == (False, True, None, 1617665) + + +def test_valid_api_v1_model_versions_url_red_domain(): + url = "https://civitai.red/api/v1/model-versions/1617665" + assert check_civitai_url(url) == (False, True, None, 1617665) + + +def test_www_subdomain_red_is_accepted(): + url = "https://www.civitai.red/models/43331?version=12345" + assert check_civitai_url(url) == (True, False, 43331, 12345) + + +def test_non_evil_civitai_red_url(): + url = "https://evilcivitai.red/models/43331?version=12345" + assert check_civitai_url(url) == (False, False, None, None) + + +def test_red_as_spoofed_subdomain_of_other_tld(): + url = "https://civitai.red.evil.com/models/43331" + assert check_civitai_url(url) == (False, False, None, None) + + def test_valid_huggingface_url(): url = "https://huggingface.co/CompVis/stable-diffusion-v1-4/resolve/main/sd-v1-4.ckpt" assert check_huggingface_url(url) == (True, "CompVis/stable-diffusion-v1-4", "sd-v1-4.ckpt", None, "main")