Coverage for src/local_deep_research/utilities/url_utils.py: 98%

1"""URL utility functions for the local deep research application.""" 

2 

3from functools import lru_cache 

4from typing import Optional 

5from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit 

6 

7from loguru import logger 

8 

9from ..security import redact_url_for_log, validate_url 

10from ..security.network_utils import is_private_ip 

11 

12# Re-export for backwards compatibility 

13__all__ = [ 

14 "normalize_url", 

15 "is_private_ip", 

16 "canonical_url_key", 

17 "is_safe_custom_llm_endpoint", 

18] 

19 

20# Tracking query parameter keys (matched lowercased). 

21_TRACKING_PARAMS = frozenset( 

22 { 

23 "fbclid", 

24 "gclid", 

25 "msclkid", 

26 "yclid", 

27 "dclid", 

28 "gad_source", 

29 "mc_eid", 

30 "mc_cid", 

31 "ref_src", 

32 "igshid", 

33 "_ga", 

34 "_gl", 

35 } 

36) 

37# Tracking param name prefixes (matched lowercased). 

38_TRACKING_PREFIXES = ("utm_",) 

39 

40 

41def normalize_url(raw_url: str) -> str: 

42 """ 

43 Normalize a URL to ensure it has a proper scheme and format. 

44 

45 Args: 

46 raw_url: The raw URL string to normalize 

47 

48 Returns: 

49 A properly formatted URL string 

50 

51 Examples: 

52 >>> normalize_url("localhost:11434") 

53 'http://localhost:11434' 

54 >>> normalize_url("https://example.com:11434") 

55 'https://example.com:11434' 

56 >>> normalize_url("http:example.com") 

57 'http://example.com' 

58 """ 

59 if not raw_url: 

60 raise ValueError("URL cannot be empty") 

61 

62 # Clean up the URL 

63 raw_url = raw_url.strip() 

64 

65 # First check if the URL already has a proper scheme 

66 if raw_url.startswith(("http://", "https://")): 

67 return raw_url 

68 

69 # Handle case where URL is malformed like "http:hostname" (missing //) 

70 if raw_url.startswith(("http:", "https:")) and not raw_url.startswith( 

71 ("http://", "https://") 

72 ): 

73 scheme = raw_url.split(":", 1)[0] 

74 rest = raw_url.split(":", 1)[1] 

75 return f"{scheme}://{rest}" 

76 

77 # Handle URLs that start with // 

78 if raw_url.startswith("//"): 

79 # Remove the // and process 

80 raw_url = raw_url[2:] 

81 

82 # At this point, we should have hostname:port or just hostname 

83 # Determine if this is localhost or an external host 

84 hostname = raw_url.split(":")[0].split("/")[0] 

85 

86 # Handle IPv6 addresses in brackets 

87 if hostname.startswith("[") and "]" in raw_url: 

88 # Extract the IPv6 address including brackets 

89 hostname = raw_url.split("]")[0] + "]" 

90 

91 # Use http for local/private addresses, https for external hosts 

92 scheme = "http" if is_private_ip(hostname) else "https" 

93 

94 return f"{scheme}://{raw_url}" 

95 

96 

97@lru_cache(maxsize=1024) 

98def canonical_url_key(url: str) -> str: 

99 """Return a canonical form of ``url`` suitable for deduplication and 

100 display in a Sources / citations listing. 

101 

102 The canonical form: 

103 - lowercases scheme and host (paths stay case-sensitive), 

104 - strips userinfo (``user:pass@`` — never leak creds), 

105 - strips default ports (80/http, 443/https), 

106 - strips fragments, 

107 - drops tracking query params (``utm_*``, ``fbclid``, ``gclid``, 

108 ``msclkid``, ``yclid``, ``dclid``, ``gad_source``, ``mc_eid``, 

109 ``mc_cid``, ``ref_src``, ``igshid``, ``_ga``, ``_gl``), 

110 - trims a trailing ``/`` from non-root paths. 

111 

112 Click-through behavior is preserved — tracking params carry no 

113 content, and mainstream browsers already strip them automatically. 

114 Percent-encoding is not normalized; query param order is preserved 

115 as-is. 

116 

117 Falls back to ``url.strip()`` when the input is not a recognizable 

118 absolute URL (e.g. ``mailto:``, ``data:``, or protocol-relative 

119 ``//host/p``), since canonicalization would be ambiguous. 

120 """ 

121 if not url: 

122 return "" 

123 try: 

124 parsed = urlsplit(url) 

125 except Exception: 

126 return url.strip() 

127 # Require both a scheme and a netloc; otherwise canonicalization is 

128 # ambiguous (mailto:, data:, protocol-relative, etc.). 

129 if not parsed.scheme or not parsed.netloc: 

130 return url.strip() 

131 

132 scheme = parsed.scheme.lower() 

133 

134 # Strip userinfo (user:pass@host) from netloc. 

135 netloc = parsed.netloc.rsplit("@", 1)[-1] 

136 

137 # Split host/port carefully so IPv6 literals survive. 

138 if netloc.startswith("["): 

139 end = netloc.find("]") 

140 host = netloc[: end + 1] 

141 rest = netloc[end + 1 :] 

142 port = rest[1:] if rest.startswith(":") else "" 

143 elif ":" in netloc: 

144 host, _, port = netloc.rpartition(":") 

145 host = host.lower() 

146 else: 

147 host, port = netloc.lower(), "" 

148 

149 if (scheme == "https" and port == "443") or ( 

150 scheme == "http" and port == "80" 

151 ): 

152 port = "" 

153 netloc = f"{host}:{port}" if port else host 

154 

155 # Filter query params case-insensitively on key; preserve order/values. 

156 if parsed.query: 

157 pairs = parse_qsl(parsed.query, keep_blank_values=True) 

158 kept = [ 

159 (k, v) 

160 for k, v in pairs 

161 if not ( 

162 k.lower() in _TRACKING_PARAMS 

163 or any(k.lower().startswith(p) for p in _TRACKING_PREFIXES) 

164 ) 

165 ] 

166 query_str = urlencode(kept, doseq=True) if kept else "" 

167 else: 

168 query_str = "" 

169 

170 path = parsed.path 

171 if path and path != "/" and path.endswith("/"): 

172 path = path.rstrip("/") 

173 

174 return urlunsplit((scheme, netloc, path, query_str, "")) 

175 

176 

177def is_safe_custom_llm_endpoint(custom_endpoint: Optional[str]) -> bool: 

178 """SSRF guard for a user-supplied custom LLM endpoint, applied at the 

179 request boundary as fail-fast defense-in-depth. 

180 

181 The endpoint is normalized exactly as the OpenAI-compatible provider 

182 normalizes it (:func:`normalize_url`), so scheme-less local endpoints 

183 such as ``localhost:11434`` or ``192.168.1.10:8000`` are handled the 

184 same way the provider handles them, then validated with 

185 :func:`validate_url` allowing private IPs / localhost. That accepts 

186 local LLM backends (Ollama / LM Studio / vLLM) while still blocking 

187 cloud-metadata and link-local targets. An empty / unset endpoint is 

188 safe (there is nothing to send to). On rejection a redacted warning 

189 is logged (the raw URL may carry credentials). 

190 

191 This is not the sole protection: the OpenAI-compatible provider's 

192 ``assert_base_url_safe`` re-validates the same URL before the 

193 LangChain client is constructed. This guard simply rejects early — 

194 before any DB row is written or research thread is spawned — and 

195 keeps the endpoint out of the logs. 

196 """ 

197 endpoint = (custom_endpoint or "").strip() 

198 if not endpoint: 

199 return True 

200 candidate = normalize_url(endpoint) 

201 if validate_url(candidate, allow_private_ips=True): 

202 return True 

203 logger.warning( 

204 "SSRF protection: rejected custom_endpoint URL: {}", 

205 redact_url_for_log(candidate), 

206 ) 

207 return False