kaiyuan/chatgpt-on-wechat
1
0
Fork 0
mirror of https://github.com/zhayujie/chatgpt-on-wechat.git synced 2026-07-17 11:07:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'upstream/master'

Browse Source
This commit is contained in:
yangziyu-hhh
2026-06-15 13:17:12 +08:00
parent 6e3933be30 e2cb9e11b0
commit e7a069b060
42 changed files with 2230 additions and 234 deletions
Download Patch File Download Diff File Expand all files Collapse all files

240
tests/test_custom_provider.py Normal file
View File

@@ -0,0 +1,240 @@
# encoding:utf-8
"""
Unit tests for multiple custom (OpenAI-compatible) provider support (issue #2838).
Covers models/custom_provider.py:
- Backward compatibility: legacy custom_api_key / custom_api_base fallback
- Multi-provider selection via bot_type "custom:<id>" routing
- parse_custom_bot_type helper
- Robustness against malformed config (missing id, non-dict, non-list)
"""
import sys
import os
import unittest
# Add project root to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
import config as config_module
from config import Config
def set_conf(d):
"""Install a fresh Config as the global config used by conf()."""
config_module.config = Config(d)
class TestParseCustomBotType(unittest.TestCase):
"""parse_custom_bot_type() parsing logic."""
def setUp(self):
from models.custom_provider import parse_custom_bot_type
self.parse = parse_custom_bot_type
def test_legacy_custom(self):
is_custom, pid = self.parse("custom")
self.assertTrue(is_custom)
self.assertEqual(pid, "")
def test_custom_with_id(self):
is_custom, pid = self.parse("custom:3f2a9c1b")
self.assertTrue(is_custom)
self.assertEqual(pid, "3f2a9c1b")
def test_non_custom(self):
is_custom, pid = self.parse("openai")
self.assertFalse(is_custom)
self.assertEqual(pid, "")
def test_empty(self):
is_custom, pid = self.parse("")
self.assertFalse(is_custom)
self.assertEqual(pid, "")
def test_none(self):
is_custom, pid = self.parse(None)
self.assertFalse(is_custom)
self.assertEqual(pid, "")
class TestResolveCustomCredentials(unittest.TestCase):
"""resolve_custom_credentials() resolution order and fallbacks."""
def setUp(self):
from models.custom_provider import resolve_custom_credentials, get_custom_providers
self.resolve = resolve_custom_credentials
self.get_providers = get_custom_providers
# --- Backward compatibility ---
def test_legacy_fallback_when_no_providers(self):
set_conf({
"bot_type": "custom",
"custom_api_key": "legacy-key",
"custom_api_base": "https://legacy.example.com/v1",
})
self.assertEqual(
self.resolve(),
("legacy-key", "https://legacy.example.com/v1", None),
)
def test_empty_config(self):
set_conf({"bot_type": "custom"})
self.assertEqual(self.resolve(), ("", None, None))
# --- Multi-provider selection via bot_type ---
def test_provider_selected_by_id(self):
set_conf({
"bot_type": "custom:abc12345",
"custom_providers": [
{"id": "sf001", "name": "siliconflow", "api_key": "sf-key",
"api_base": "https://api.siliconflow.cn/v1", "model": "deepseek-ai/DeepSeek-V3"},
{"id": "abc12345", "name": "qiniu", "api_key": "qn-key",
"api_base": "https://api.qnaigc.com/v1", "model": "deepseek-v3"},
],
})
self.assertEqual(
self.resolve(),
("qn-key", "https://api.qnaigc.com/v1", "deepseek-v3"),
)
def test_id_not_found_falls_back_to_legacy(self):
set_conf({
"bot_type": "custom:ghost",
"custom_api_key": "legacy-key",
"custom_api_base": "https://legacy.example.com/v1",
"custom_providers": [
{"id": "sf001", "name": "siliconflow", "api_key": "sf-key",
"api_base": "https://api.siliconflow.cn/v1"},
],
})
self.assertEqual(
self.resolve(),
("legacy-key", "https://legacy.example.com/v1", None),
)
def test_provider_without_model_returns_none_model(self):
set_conf({
"bot_type": "custom:local01",
"custom_providers": [
{"id": "local01", "name": "local", "api_key": "", "api_base": "http://localhost:11434/v1"},
],
})
self.assertEqual(
self.resolve(),
("", "http://localhost:11434/v1", None),
)
# --- Robustness against malformed config ---
def test_malformed_entries_filtered_and_fallback(self):
set_conf({
"bot_type": "custom:nope",
"custom_api_key": "legacy-key",
"custom_api_base": "https://legacy.example.com/v1",
"custom_providers": [
{"name": "no-id", "api_key": "no-id-key"}, # invalid: no id
"not-a-dict", # invalid: wrong type
],
})
# All entries invalid -> treated as empty -> legacy fallback
self.assertEqual(
self.resolve(),
("legacy-key", "https://legacy.example.com/v1", None),
)
def test_get_custom_providers_filters_invalid(self):
set_conf({
"bot_type": "custom",
"custom_providers": [
{"id": "ok1", "name": "ok", "api_key": "k", "api_base": "https://x/v1"},
{"name": "no-id", "api_key": "no-id"}, # dropped: no id
123, # dropped
],
})
providers = self.get_providers()
self.assertEqual(len(providers), 1)
self.assertEqual(providers[0]["id"], "ok1")
def test_custom_providers_not_a_list_falls_back(self):
set_conf({
"bot_type": "custom",
"custom_api_key": "legacy-key",
"custom_api_base": "https://legacy.example.com/v1",
"custom_providers": "oops-a-string",
})
self.assertEqual(
self.resolve(),
("legacy-key", "https://legacy.example.com/v1", None),
)
class TestGenerateProviderId(unittest.TestCase):
"""generate_provider_id() produces valid short ids."""
def test_length_and_hex(self):
from models.custom_provider import generate_provider_id
pid = generate_provider_id()
self.assertEqual(len(pid), 8)
# Must be valid hex characters
int(pid, 16)
def test_uniqueness(self):
from models.custom_provider import generate_provider_id
ids = {generate_provider_id() for _ in range(100)}
self.assertEqual(len(ids), 100)
class TestConfigDefaults(unittest.TestCase):
"""The new config fields must exist with safe defaults."""
def test_default_config_has_custom_providers(self):
from config import available_setting
self.assertIn("custom_providers", available_setting)
self.assertEqual(available_setting["custom_providers"], [])
def test_default_config_no_custom_active_provider(self):
"""custom_active_provider was removed — replaced by bot_type routing."""
from config import available_setting
self.assertNotIn("custom_active_provider", available_setting)
class TestDragSensitiveNested(unittest.TestCase):
"""drag_sensitive() must mask api_key in nested structures."""
def test_nested_api_key_masked(self):
from config import drag_sensitive
import json
test_config = {
"open_ai_api_key": "sk-1234567890abcdef",
"custom_providers": [
{"id": "x1", "name": "test", "api_key": "sk-nested-secret-key-long", "api_base": "https://x/v1"}
],
}
result = drag_sensitive(test_config)
# Top-level key should be masked
self.assertNotIn("1234567890abcdef", str(result))
# Nested key should also be masked
self.assertNotIn("nested-secret-key-long", str(result))
# But the id/name/api_base should not be masked
self.assertIn("x1", str(result))
self.assertIn("test", str(result))
self.assertIn("https://x/v1", str(result))
def test_string_config_masked(self):
from config import drag_sensitive
import json
test_str = json.dumps({
"open_ai_api_key": "sk-1234567890abcdef",
"custom_providers": [
{"id": "x1", "api_key": "sk-nested-very-long-secret"}
],
})
result = drag_sensitive(test_str)
self.assertNotIn("1234567890abcdef", result)
self.assertNotIn("nested-very-long-secret", result)
if __name__ == "__main__":
unittest.main()

289
tests/test_custom_provider_handlers.py Normal file
View File

@@ -0,0 +1,289 @@
# encoding:utf-8
"""
Unit tests for the multi custom-provider management API (issue #2838, web UI).
Covers channel/web/web_channel.py::ModelsHandler:
- _custom_provider_cards / _provider_overview expansion
- _handle_set_custom_provider (create / edit / activate)
- _handle_delete_custom_provider
- _handle_set_active_custom_provider
Uses id-based routing (bot_type: "custom:<id>") — no custom_active_provider.
"""
import json
import os
import sys
import types
import unittest
# Add project root to path.
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
# Stub the web.py framework so web_channel imports without the dependency.
if "web" not in sys.modules:
_web_stub = types.ModuleType("web")
_web_stub.header = lambda *a, **k: None
_web_stub.data = lambda: b"{}"
_web_stub.ctx = types.SimpleNamespace()
sys.modules["web"] = _web_stub
import config as config_module
from config import Config
from channel.web.web_channel import ModelsHandler
def set_conf(d):
"""Install a fresh Config as the global config used by conf()."""
config_module.config = Config(d)
class _HandlerHarness:
"""Test double around ModelsHandler that captures persisted config in
memory instead of touching config.json, and no-ops the Bridge reset."""
def __init__(self):
self.handler = ModelsHandler.__new__(ModelsHandler)
self._file_cfg = {}
self.bridge_resets = 0
# Patch the disk + bridge boundary on this instance.
self.handler._read_file_config = lambda: dict(self._file_cfg)
self.handler._write_file_config = self._capture_write
self.handler._reset_bridge = self._capture_reset
def _capture_write(self, data):
self._file_cfg = dict(data)
def _capture_reset(self):
self.bridge_resets += 1
def call(self, **payload):
# Resolve the bound method by action for convenience.
action = payload.get("action")
method = {
"set_custom_provider": self.handler._handle_set_custom_provider,
"delete_custom_provider": self.handler._handle_delete_custom_provider,
"set_active_custom_provider": self.handler._handle_set_active_custom_provider,
}[action]
return json.loads(method(payload))
class TestSetCustomProvider(unittest.TestCase):
def setUp(self):
set_conf({"bot_type": "custom", "custom_providers": []})
self.h = _HandlerHarness()
def test_create_provider_does_not_hijack_bot_type(self):
"""Creating a provider without make_active must not change bot_type."""
res = self.h.call(action="set_custom_provider", name="siliconflow",
api_base="https://api.siliconflow.cn/v1", api_key="sf-key")
self.assertEqual(res["status"], "success")
self.assertTrue(res["created"])
self.assertIn("id", res)
# bot_type must remain unchanged — no auto-activation.
bot_type = config_module.conf().get("bot_type")
self.assertEqual(bot_type, "custom") # unchanged from setUp
providers = config_module.conf().get("custom_providers")
self.assertEqual(len(providers), 1)
self.assertEqual(providers[0]["id"], res["id"])
self.assertEqual(providers[0]["name"], "siliconflow")
self.assertEqual(self.h.bridge_resets, 1)
def test_create_with_make_active_switches_bot_type(self):
"""Creating a provider with make_active=true must switch bot_type."""
res = self.h.call(action="set_custom_provider", name="siliconflow",
api_base="https://api.siliconflow.cn/v1", api_key="sf-key",
make_active=True)
self.assertEqual(res["status"], "success")
bot_type = config_module.conf().get("bot_type")
self.assertEqual(bot_type, f"custom:{res['id']}")
def test_create_requires_api_base(self):
res = self.h.call(action="set_custom_provider", name="x", api_key="k")
self.assertEqual(res["status"], "error")
self.assertIn("api_base", res["message"])
def test_create_requires_name(self):
res = self.h.call(action="set_custom_provider", name="", api_base="https://x/v1")
self.assertEqual(res["status"], "error")
def test_second_provider_does_not_steal_active(self):
# Explicitly activate the first provider.
res1 = self.h.call(action="set_custom_provider", name="a",
api_base="https://a/v1", api_key="ak", make_active=True)
res2 = self.h.call(action="set_custom_provider", name="b",
api_base="https://b/v1", api_key="bk")
self.assertTrue(res2["created"])
# First provider stays active — second creation doesn't steal it.
bot_type = config_module.conf().get("bot_type")
self.assertEqual(bot_type, f"custom:{res1['id']}")
def test_make_active_flag(self):
self.h.call(action="set_custom_provider", name="a",
api_base="https://a/v1", api_key="ak")
res2 = self.h.call(action="set_custom_provider", name="b",
api_base="https://b/v1", api_key="bk", make_active=True)
bot_type = config_module.conf().get("bot_type")
self.assertEqual(bot_type, f"custom:{res2['id']}")
def test_edit_keeps_key_when_omitted(self):
res = self.h.call(action="set_custom_provider", name="a",
api_base="https://a/v1", api_key="secret")
pid = res["id"]
# Edit only the base; omit api_key.
res2 = self.h.call(action="set_custom_provider", name="a",
id=pid, api_base="https://a2/v1")
self.assertEqual(res2["status"], "success")
self.assertFalse(res2["created"])
providers = config_module.conf().get("custom_providers")
self.assertEqual(providers[0]["api_base"], "https://a2/v1")
self.assertEqual(providers[0]["api_key"], "secret") # preserved
def test_edit_can_rename(self):
res = self.h.call(action="set_custom_provider", name="old",
api_base="https://a/v1", api_key="ak")
pid = res["id"]
res2 = self.h.call(action="set_custom_provider", name="new",
id=pid, api_base="https://a/v1")
self.assertEqual(res2["status"], "success")
providers = config_module.conf().get("custom_providers")
self.assertEqual(providers[0]["name"], "new")
# ID stays the same
self.assertEqual(providers[0]["id"], pid)
def test_edit_clears_model_when_empty(self):
res = self.h.call(action="set_custom_provider", name="a",
api_base="https://a/v1", api_key="ak", model="m1")
pid = res["id"]
self.assertEqual(config_module.conf().get("custom_providers")[0]["model"], "m1")
self.h.call(action="set_custom_provider", name="a", id=pid,
api_base="https://a/v1", model="")
self.assertNotIn("model", config_module.conf().get("custom_providers")[0])
class TestDeleteCustomProvider(unittest.TestCase):
def setUp(self):
set_conf({"bot_type": "custom", "custom_providers": []})
self.h = _HandlerHarness()
self.res_a = self.h.call(action="set_custom_provider", name="a",
api_base="https://a/v1", api_key="ak",
make_active=True)
self.res_b = self.h.call(action="set_custom_provider", name="b",
api_base="https://b/v1", api_key="bk")
def test_delete_unknown(self):
res = self.h.call(action="delete_custom_provider", id="ghost")
self.assertEqual(res["status"], "error")
def test_delete_non_active(self):
res = self.h.call(action="delete_custom_provider", id=self.res_b["id"])
self.assertEqual(res["status"], "success")
ids = [p["id"] for p in config_module.conf().get("custom_providers")]
self.assertEqual(ids, [self.res_a["id"]])
# bot_type unchanged (still pointing to a)
self.assertEqual(config_module.conf().get("bot_type"), f"custom:{self.res_a['id']}")
def test_delete_active_falls_back_to_first_remaining(self):
# 'a' is active (created first); deleting it should re-point to 'b'.
self.assertEqual(config_module.conf().get("bot_type"), f"custom:{self.res_a['id']}")
res = self.h.call(action="delete_custom_provider", id=self.res_a["id"])
self.assertEqual(res["status"], "success")
self.assertEqual(config_module.conf().get("bot_type"), f"custom:{self.res_b['id']}")
def test_delete_last_reverts_to_legacy(self):
self.h.call(action="delete_custom_provider", id=self.res_a["id"])
self.h.call(action="delete_custom_provider", id=self.res_b["id"])
self.assertEqual(config_module.conf().get("custom_providers"), [])
# When all providers deleted, reverts to legacy "custom"
self.assertEqual(config_module.conf().get("bot_type"), "custom")
class TestSetActiveCustomProvider(unittest.TestCase):
def setUp(self):
set_conf({"bot_type": "custom", "custom_providers": []})
self.h = _HandlerHarness()
self.res_a = self.h.call(action="set_custom_provider", name="a",
api_base="https://a/v1", api_key="ak",
make_active=True)
self.res_b = self.h.call(action="set_custom_provider", name="b",
api_base="https://b/v1", api_key="bk")
def test_set_active_valid(self):
res = self.h.call(action="set_active_custom_provider", id=self.res_b["id"])
self.assertEqual(res["status"], "success")
self.assertEqual(config_module.conf().get("bot_type"), f"custom:{self.res_b['id']}")
def test_set_active_unknown(self):
res = self.h.call(action="set_active_custom_provider", id="ghost")
self.assertEqual(res["status"], "error")
# bot_type unchanged
self.assertEqual(config_module.conf().get("bot_type"), f"custom:{self.res_a['id']}")
def test_activation_syncs_model_to_global(self):
"""Activating a provider must write its model into global model field."""
set_conf({"bot_type": "custom", "custom_providers": [], "model": "gpt-4o"})
h = _HandlerHarness()
res = h.call(action="set_custom_provider", name="sf",
api_base="https://sf/v1", api_key="k", model="deepseek-v3",
make_active=True)
# Global model field should now be the provider's model.
self.assertEqual(config_module.conf().get("model"), "deepseek-v3")
self.assertEqual(config_module.conf().get("bot_type"), f"custom:{res['id']}")
def test_activation_without_model_keeps_global_model(self):
"""Activating a provider with no model must not overwrite global model."""
set_conf({"bot_type": "custom", "custom_providers": [], "model": "gpt-4o"})
h = _HandlerHarness()
h.call(action="set_custom_provider", name="local",
api_base="http://localhost:11434/v1", api_key="",
make_active=True)
# Global model field should remain unchanged.
self.assertEqual(config_module.conf().get("model"), "gpt-4o")
class TestProviderOverviewExpansion(unittest.TestCase):
"""_provider_overview / _custom_provider_cards should expand the list."""
def test_no_custom_providers_keeps_single_card(self):
set_conf({"bot_type": "custom", "custom_providers": []})
cards = ModelsHandler._custom_provider_cards(config_module.conf())
self.assertEqual(cards, [])
overview = ModelsHandler._provider_overview()
custom_cards = [c for c in overview if c.get("id") == "custom"]
# Legacy single custom card remains present.
self.assertEqual(len(custom_cards), 1)
def test_multi_providers_expand_into_cards(self):
set_conf({
"bot_type": "custom:id_b",
"custom_providers": [
{"id": "id_a", "name": "a", "api_key": "ak", "api_base": "https://a/v1"},
{"id": "id_b", "name": "b", "api_key": "bk", "api_base": "https://b/v1", "model": "m"},
],
})
overview = ModelsHandler._provider_overview()
custom_cards = [c for c in overview if c.get("is_custom")]
self.assertEqual(len(custom_cards), 2)
by_id = {c["custom_id"]: c for c in custom_cards}
self.assertEqual(by_id["id_a"]["id"], "custom:id_a")
self.assertFalse(by_id["id_a"]["active"])
self.assertTrue(by_id["id_b"]["active"])
self.assertEqual(by_id["id_b"]["model"], "m")
# No single legacy "custom" card when expanded.
self.assertFalse(any(c.get("id") == "custom" for c in overview))
def test_no_active_shows_none_active(self):
"""When bot_type is plain 'custom', no card is marked active."""
set_conf({
"bot_type": "custom",
"custom_providers": [
{"id": "id_a", "name": "a", "api_key": "ak", "api_base": "https://a/v1"},
{"id": "id_b", "name": "b", "api_key": "bk", "api_base": "https://b/v1"},
],
})
cards = ModelsHandler._custom_provider_cards(config_module.conf())
active_cards = [c for c in cards if c.get("active")]
self.assertEqual(len(active_cards), 0)
if __name__ == "__main__":
unittest.main()

121
tests/test_robustness_fixes.py Normal file
View File

@@ -0,0 +1,121 @@
# encoding:utf-8
"""
Unit tests for robustness fixes:
1. ChatChannel.cancel_session / cancel_all_session must not raise KeyError
when a session has been produced but no task has been dispatched yet
(so self.futures[session_id] does not exist).
2. common.utils.compress_imgfile must terminate (no infinite loop / invalid
PIL quality) when an image cannot be compressed below max_size.
"""
import io
import os
import sys
import types
import unittest
from unittest.mock import MagicMock
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
# =============================================================================
# 1. cancel_session / cancel_all_session KeyError regression
# =============================================================================
class TestCancelSessionMissingFutures(unittest.TestCase):
"""A session may exist in self.sessions before any future is recorded."""
def _make_channel(self):
# Import lazily and build a bare object without running __init__,
# to avoid pulling the full channel setup / config.
from channel.chat_channel import ChatChannel
ch = ChatChannel.__new__(ChatChannel)
import threading
ch.lock = threading.RLock()
# A produced session whose future has NOT been dispatched yet.
queue = MagicMock()
queue.qsize.return_value = 0
semaphore = MagicMock()
ch.sessions = {"sid": [queue, semaphore]}
ch.futures = {} # intentionally empty: consume() never ran
return ch
def test_cancel_session_no_futures_entry(self):
ch = self._make_channel()
# Should not raise KeyError.
try:
ch.cancel_session("sid")
except KeyError:
self.fail("cancel_session raised KeyError when futures entry missing")
def test_cancel_all_session_no_futures_entry(self):
ch = self._make_channel()
try:
ch.cancel_all_session()
except KeyError:
self.fail("cancel_all_session raised KeyError when futures entry missing")
def test_cancel_session_cancels_existing_futures(self):
ch = self._make_channel()
fut = MagicMock()
ch.futures["sid"] = [fut]
ch.cancel_session("sid")
fut.cancel.assert_called_once()
# =============================================================================
# 2. compress_imgfile termination
# =============================================================================
class TestCompressImgfileTermination(unittest.TestCase):
"""compress_imgfile must always return, even for incompressible input."""
def setUp(self):
# Skip if Pillow is not available in the test environment.
try:
import PIL # noqa: F401
except ImportError:
self.skipTest("Pillow not installed")
def _make_image_buf(self, size=(64, 64)):
from PIL import Image
import random
img = Image.new("RGB", size)
# Fill with random noise so JPEG cannot compress it well.
pixels = [
(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
for _ in range(size[0] * size[1])
]
img.putdata(pixels)
buf = io.BytesIO()
img.save(buf, "JPEG", quality=95)
buf.seek(0)
return buf
def test_returns_when_target_unreachable(self):
from common.utils import compress_imgfile
buf = self._make_image_buf()
# An impossibly small target that even quality=10 won't reach.
out = compress_imgfile(buf, max_size=10)
self.assertIsInstance(out, io.BytesIO)
# Verify the result is still a valid JPEG (PIL never got invalid quality).
from PIL import Image
out.seek(0)
img = Image.open(out)
img.verify()
def test_no_compression_needed_returns_same_object(self):
from common.utils import compress_imgfile
buf = self._make_image_buf()
size = buf.getbuffer().nbytes
out = compress_imgfile(buf, max_size=size + 1)
self.assertIs(out, buf)
if __name__ == "__main__":
unittest.main()

194
tests/test_security_ssrf_path_traversal.py Normal file
View File

@@ -0,0 +1,194 @@
# encoding:utf-8
"""
Unit tests for security fixes:
1. Vision tool SSRF protection (issue #2878, #2872)
2. Skill service path traversal protection (issue #2873)
"""
import os
import sys
import tempfile
import types
import unittest
from unittest.mock import patch, MagicMock
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
# Stub 'requests' if not installed so vision.py can be imported for testing.
if "requests" not in sys.modules:
_requests_stub = types.ModuleType("requests")
_requests_stub.get = lambda *a, **k: None
sys.modules["requests"] = _requests_stub
# =============================================================================
# Vision SSRF tests
# =============================================================================
class TestVisionSSRFValidation(unittest.TestCase):
"""Test that _validate_url_safe blocks internal/private URLs."""
def setUp(self):
from agent.tools.vision.vision import Vision
self.validate = Vision._validate_url_safe
def test_loopback_ipv4_blocked(self):
"""127.0.0.1 must be rejected."""
with self.assertRaises(ValueError) as ctx:
self.validate("http://127.0.0.1/canary.png")
self.assertIn("non-public", str(ctx.exception))
def test_loopback_localhost_blocked(self):
"""localhost must be rejected."""
with self.assertRaises(ValueError) as ctx:
self.validate("http://localhost/canary.png")
self.assertIn("non-public", str(ctx.exception))
def test_private_10_network_blocked(self):
"""10.x.x.x RFC1918 must be rejected."""
with patch("socket.getaddrinfo") as mock_gai:
mock_gai.return_value = [
(2, 1, 6, "", ("10.0.0.1", 0)),
]
with self.assertRaises(ValueError) as ctx:
self.validate("http://internal.corp/image.png")
self.assertIn("non-public", str(ctx.exception))
def test_private_172_network_blocked(self):
"""172.16.x.x RFC1918 must be rejected."""
with patch("socket.getaddrinfo") as mock_gai:
mock_gai.return_value = [
(2, 1, 6, "", ("172.16.0.1", 0)),
]
with self.assertRaises(ValueError) as ctx:
self.validate("http://internal.corp/image.png")
self.assertIn("non-public", str(ctx.exception))
def test_private_192_168_blocked(self):
"""192.168.x.x RFC1918 must be rejected."""
with patch("socket.getaddrinfo") as mock_gai:
mock_gai.return_value = [
(2, 1, 6, "", ("192.168.1.1", 0)),
]
with self.assertRaises(ValueError) as ctx:
self.validate("http://router.local/image.png")
self.assertIn("non-public", str(ctx.exception))
def test_link_local_blocked(self):
"""169.254.x.x (link-local / cloud metadata) must be rejected."""
with patch("socket.getaddrinfo") as mock_gai:
mock_gai.return_value = [
(2, 1, 6, "", ("169.254.169.254", 0)),
]
with self.assertRaises(ValueError) as ctx:
self.validate("http://metadata.google.internal/image.png")
self.assertIn("non-public", str(ctx.exception))
def test_ipv6_loopback_blocked(self):
"""::1 (IPv6 loopback) must be rejected."""
with patch("socket.getaddrinfo") as mock_gai:
mock_gai.return_value = [
(10, 1, 6, "", ("::1", 0, 0, 0)),
]
with self.assertRaises(ValueError) as ctx:
self.validate("http://[::1]/image.png")
self.assertIn("non-public", str(ctx.exception))
def test_public_url_allowed(self):
"""A URL resolving to a public IP should pass validation."""
with patch("socket.getaddrinfo") as mock_gai:
mock_gai.return_value = [
(2, 1, 6, "", ("151.101.1.140", 0)),
]
# Should not raise
self.validate("https://cdn.example.com/image.png")
def test_no_hostname_rejected(self):
"""A URL with no host must be rejected."""
with self.assertRaises(ValueError) as ctx:
self.validate("http:///path/to/image.png")
self.assertIn("no hostname", str(ctx.exception))
def test_non_http_scheme_rejected(self):
"""file:// and ftp:// schemes must be rejected."""
with self.assertRaises(ValueError) as ctx:
self.validate("file:///etc/passwd")
self.assertIn("scheme", str(ctx.exception))
def test_dns_failure_rejected(self):
"""Unresolvable hostname must be rejected."""
import socket as sock_mod
with patch("socket.getaddrinfo", side_effect=sock_mod.gaierror("Name does not resolve")):
with self.assertRaises(ValueError) as ctx:
self.validate("http://nonexistent.invalid/img.png")
self.assertIn("Cannot resolve", str(ctx.exception))
# =============================================================================
# Skill service path traversal tests
# =============================================================================
class TestSkillServicePathTraversal(unittest.TestCase):
"""Test that _safe_skill_dir blocks path traversal attempts."""
def setUp(self):
self.tmp_root = tempfile.mkdtemp()
# Create a minimal SkillManager mock with custom_dir set.
from agent.skills.service import SkillService
mock_manager = MagicMock()
mock_manager.custom_dir = self.tmp_root
self.svc = SkillService(mock_manager)
def tearDown(self):
import shutil
shutil.rmtree(self.tmp_root, ignore_errors=True)
def test_normal_name_allowed(self):
"""A simple name like 'my-skill' should produce a valid path."""
result = self.svc._safe_skill_dir("my-skill")
expected = os.path.realpath(os.path.join(self.tmp_root, "my-skill"))
self.assertEqual(result, expected)
def test_dotdot_traversal_blocked(self):
"""'../escaped' must be rejected."""
with self.assertRaises(ValueError) as ctx:
self.svc._safe_skill_dir("../escaped")
self.assertIn("path traversal", str(ctx.exception))
def test_nested_dotdot_blocked(self):
"""'foo/../../escaped' must be rejected."""
with self.assertRaises(ValueError) as ctx:
self.svc._safe_skill_dir("foo/../../escaped")
self.assertIn("path traversal", str(ctx.exception))
def test_absolute_path_blocked(self):
"""'/tmp/evil' must be rejected."""
with self.assertRaises(ValueError) as ctx:
self.svc._safe_skill_dir("/tmp/evil")
self.assertIn("path traversal", str(ctx.exception))
def test_backslash_path_blocked(self):
r"""'\\server\share' must be rejected."""
with self.assertRaises(ValueError) as ctx:
self.svc._safe_skill_dir("\\server\\share")
self.assertIn("path traversal", str(ctx.exception))
def test_empty_name_blocked(self):
"""Empty name must be rejected."""
with self.assertRaises(ValueError):
self.svc._safe_skill_dir("")
def test_whitespace_only_blocked(self):
"""Whitespace-only name must be rejected."""
with self.assertRaises(ValueError):
self.svc._safe_skill_dir(" ")
def test_subdir_name_allowed(self):
"""A name with a forward slash but no traversal is allowed if it stays in root."""
# e.g. "category/skill-name" is a valid nested skill directory
result = self.svc._safe_skill_dir("category/skill-name")
expected = os.path.realpath(os.path.join(self.tmp_root, "category/skill-name"))
self.assertEqual(result, expected)
if __name__ == "__main__":
unittest.main()