fix: accept civitai.red as an allowed CivitAI host (#436)

Signed-off-by: Alexander Piskun <bigcat88@icloud.com>

Author: bigcat88

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)

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"
 

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")