Add security check to see if a URL is pointing to internal IP

src/agentscope/manager/_file.py

@@ -34,21 +34,19 @@ def _get_text_embedding_record_hash(
  if isinstance(embedding_model, dict):
  # Format the dict to avoid duplicate keys
  embedding_model = json.dumps(embedding_model, sort_keys=True)
- elif isinstance(embedding_model, str):
- embedding_model_hash = _hash_string(embedding_model, hash_method)
- else:
+ elif not isinstance(embedding_model, str):
  raise RuntimeError(
  f"The embedding model must be a string or a dict, got "
  f"{type(embedding_model)}.",
  )
+ embedding_model_hash = _hash_string(embedding_model, hash_method)
 
  # Calculate the embedding id by hashing the hash codes of the
  # original data and the embedding model
  record_hash = _hash_string(
  original_data_hash + embedding_model_hash,
  hash_method,
  )
-
  return record_hash
 
 

src/agentscope/rag/llama_index_knowledge.py

@@ -203,8 +203,9 @@ def __init__(
  )
 
  if persist_root is None:
- persist_root = FileManager.get_instance().run_dir or "./"
+ persist_root = FileManager.get_instance().cache_dir or "./"
  self.persist_dir = os.path.join(persist_root, knowledge_id)
+ logger.info(f"** persist_dir: {self.persist_dir}")
  self.emb_model = emb_model
  self.overwrite_index = overwrite_index
  self.showprogress = showprogress

src/agentscope/service/web/web_digest.py

@@ -3,6 +3,8 @@
 import json
 from urllib.parse import urlparse
 from typing import Optional, Callable, Sequence, Any
+import socket
+import ipaddress
 import requests
 from loguru import logger
 
@@ -37,12 +39,46 @@ def is_valid_url(url: str) -> bool:
  return False # A ValueError indicates that the URL is not valid.
 
 
+def is_internal_ip_address(url: str) -> bool:
+ """
+ Check if a URL is to interal IP addresses
+ Args:
+ url (str): url to be checked
+
+ Returns:
+ bool: True if url is not to interal IP addresses,
+ False otherwise
+ """
+ parsed_url = urlparse(url)
+ hostname = parsed_url.hostname
+ if hostname is None:
+ # illegal hostname is ignore in this function
+ return False
+
+ # Resolve the hostname to an IP address
+ ip = socket.gethostbyname(hostname)
+ # Check if it's localhost or within the loopback range
+ if (
+ ip.startswith("127.")
+ or ip == "::1"
+ or ipaddress.ip_address(ip).is_private
+ ):
+ logger.warning(
+ f"Access to this URL {url} is "
+ f"restricted because it is private",
+ )
+ return True
+
+ return False
+
+
 def load_web(
  url: str,
  keep_raw: bool = True,
  html_selected_tags: Optional[Sequence[str]] = None,
  self_parse_func: Optional[Callable[[requests.Response], Any]] = None,
  timeout: int = 5,
+ exclude_internal_ips: bool = True,
 ) -> ServiceResponse:
  """Function for parsing and digesting the web page.
 
@@ -62,6 +98,8 @@ def load_web(
  The result is stored with `self_define_func`
  key
  timeout (int): timeout parameter for requests.
+ exclude_internal_ips (bool):
+ whether prevent the function access internal_ips
 
  Returns:
  `ServiceResponse`: If successful, `ServiceResponse` object is returned
@@ -87,6 +125,13 @@ def load_web(
  "selected_tags_text": xxxxx
  }
  """
+ if exclude_internal_ips and is_internal_ip_address(url):
+ return ServiceResponse(
+ ServiceExecStatus.ERROR,
+ content=f"Access to this URL {url} is restricted "
+ f"because it is private",
+ )
+
  header = {
  "Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6",
  "Cache-Control": "max-age=0",

tests/web_digest_test.py

@@ -58,7 +58,7 @@ def test_web_load(self, mock_get: MagicMock) -> None:
  mock_get.return_value = mock_response
 
  # set parameters
- fake_url = "fake-url"
+ fake_url = "http://fake-url.com"
 
  results = load_web(
  url=fake_url,
@@ -100,6 +100,12 @@ def format(
  expected_result,
  )
 
+ def test_block_internal_ips(self) -> None:
+ """test whether can prevent internal_url successfully"""
+ internal_url = "http://localhost:8080/some/path"
+ response = load_web(internal_url)
+ self.assertEqual(ServiceExecStatus.ERROR, response.status)
+
 
 # This allows the tests to be run from the command line
 if __name__ == "__main__":