refactor: Updates ZIA URL Categories to support latest API endpoints

docs: Updates ZIA URL Categories docs
tests: Updates ZIA URL Categories test suite

docsrc/zs/zia/url_categories.rst

@@ -5,6 +5,16 @@ The following methods allow for interaction with the ZIA URL Categories API endp
 
 Methods are accessible via ``zia.url_categories``
 
+
+**Guidelines for adding / updating URLs**
+
+
+- The URL must use a standard URI format.
+- The URL length cannot exceed 1024 characters.
+- The URL cannot contain non-ASCII characters.
+- The domain name before the colon (:) cannot exceed 255 characters.
+- The domain name between periods (.) cannot exceed 63 characters.
+
 .. _zia-url_categories:
 
 .. automodule:: pyzscaler.zia.url_categories

pyzscaler/utils.py

@@ -18,6 +18,12 @@ def snake_to_camel(name: str):
     return ret
 
 
+def chunker(lst, n):
+    """Yield successive n-sized chunks from lst."""
+    for i in range(0, len(lst), n):
+        yield lst[i : i + n]
+
+
 # Recursive function to convert all keys and nested keys from snake case
 # to camel case.
 def convert_keys(data):

pyzscaler/zia/url_categories.py

@@ -1,7 +1,9 @@
+import time
+
 from box import Box, BoxList
 from restfly.endpoint import APIEndpoint
 
-from pyzscaler.utils import convert_keys, snake_to_camel
+from pyzscaler.utils import chunker, convert_keys, snake_to_camel
 
 
 class URLCategoriesAPI(APIEndpoint):
@@ -20,17 +22,29 @@ def lookup(self, urls: list) -> BoxList:
             >>> zia.url_categories.lookup(['example.com', 'test.com'])
 
         """
-        payload = urls
 
-        return self._post("urlLookup", json=payload)
+        # ZIA limits each API call to 100 URLs at a rate of 1 API call per second. pyZscaler simplifies this by allowing
+        # users to submit any number of URLs and handle the chunking of the API calls on their behalf.
+        if len(urls) > 100:
+            results = BoxList()
+            for chunk in chunker(urls, 100):
+                results.extend(self._post("urlLookup", json=chunk))
+                time.sleep(1)
+            return results
+
+        else:
+            payload = urls
+            return self._post("urlLookup", json=payload)
 
-    def list_categories(self, custom_only: bool = False) -> BoxList:
+    def list_categories(self, custom_only: bool = False, only_counts: bool = False) -> BoxList:
         """
         Returns information on URL categories.
 
         Args:
             custom_only (bool):
                 Returns only custom categories if True.
+            only_counts (bool):
+                Returns only URL and keyword counts if True.
 
         Returns:
             :obj:`BoxList`: A list of information for all or custom URL categories.
@@ -45,8 +59,12 @@ def list_categories(self, custom_only: bool = False) -> BoxList:
             >>> zia.url_categories.list_categories(custom_only=True)
 
         """
+        payload = {
+            "customOnly": custom_only,
+            "includeOnlyUrlKeywordCounts": only_counts,
+        }
 
-        return self._get(f"urlCategories?customOnly={custom_only}")
+        return self._get("urlCategories", params=payload)
 
     def get_quota(self) -> Box:
         """
@@ -100,6 +118,14 @@ def add_url_category(self, name: str, super_category: str, urls: list, **kwargs)
                 Description of the category.
             custom_category (bool):
                 Set to true for custom URL category. Up to 48 custom URL categories can be added per organisation.
+            ip_ranges (list):
+                Custom IP addpress ranges associated to a URL category. This feature must be enabled on your tenancy.
+            ip_ranges_retaining_parent_category (list):
+                The retaining parent custom IP addess ranges associated to a URL category.
+            keywords (list):
+                Custom keywords associated to a URL category.
+            keywords_retaining_parent_category (list):
+                Retained custom keywords from the parent URL category that are associated with a URL category.
 
         Returns:
             :obj:`Box`: The newly configured custom URL category resource record.
@@ -110,7 +136,15 @@ def add_url_category(self, name: str, super_category: str, urls: list, **kwargs)
             >>> zia.url_categories.add_url_category(name='Beer',
             ...    super_category='ALCOHOL_TOBACCO',
             ...    urls=['xxxx.com.au', 'carltondraught.com.au'],
-            ...    description="Beers that don't taste good")
+            ...    description="Beers that don't taste good.")
+
+            Add a new category with IP ranges:
+
+            >>> zia.url_categories.add_url_category(name='Beer',
+            ...    super_category='FINANCE',
+            ...    urls=['finance.google.com'],
+            ...    description="Google Finance.",
+            ...    ip_ranges=['10.0.0.0/24'])
 
         """
 
@@ -125,6 +159,8 @@ def add_url_category(self, name: str, super_category: str, urls: list, **kwargs)
         for key, value in kwargs.items():
             payload[snake_to_camel(key)] = value
 
+        print(payload)
+
         return self._post("urlCategories", json=payload)
 
     def add_tld_category(self, name: str, tlds: list, **kwargs) -> Box:
@@ -187,6 +223,14 @@ def update_url_category(self, category_id: str, **kwargs) -> Box:
                 URLs entered will be covered by policies that reference the parent category, in addition to this one.
             description (str):
                 Description of the category.
+            ip_ranges (list):
+                Custom IP addpress ranges associated to a URL category. This feature must be enabled on your tenancy.
+            ip_ranges_retaining_parent_category (list):
+                The retaining parent custom IP addess ranges associated to a URL category.
+            keywords (list):
+                Custom keywords associated to a URL category.
+            keywords_retaining_parent_category (list):
+                Retained custom keywords from the parent URL category that are associated with a URL category.
 
         Returns:
             :obj:`Box`: The updated URL category resource record.

tests/zia/test_url_categories.py

@@ -72,6 +72,15 @@ def fixture_custom_url_categories():
     ]
 
 
+@pytest.fixture(name="url_lookups")
+def fixture_url_lookups():
+    # Generate a list of URLs for the given quantity
+    def _method(num):
+        return [f"www.{x}.com" for x in range(num)]
+
+    return _method
+
+
 @responses.activate
 def test_url_category_lookup(zia):
     lookup_response = [
@@ -93,11 +102,40 @@ def test_url_category_lookup(zia):
     assert resp[0].url == "github.com"
 
 
+@responses.activate
+def test_url_category_lookup_chunked(zia, url_lookups):
+    urls = url_lookups(250)
+
+    responses.add(
+        method="POST",
+        url="https://zsapi.zscaler.net/api/v1/urlLookup",
+        json=urls[:101],
+        status=200,
+    )
+
+    responses.add(
+        method="POST",
+        url="https://zsapi.zscaler.net/api/v1/urlLookup",
+        json=urls[101:201],
+        status=200,
+    )
+    responses.add(
+        method="POST",
+        url="https://zsapi.zscaler.net/api/v1/urlLookup",
+        json=urls[201:],
+        status=200,
+    )
+
+    resp = zia.url_categories.lookup(urls)
+    assert isinstance(resp, BoxList)
+    assert len(resp) == 250
+
+
 @responses.activate
 def test_list_categories(zia, url_categories):
     responses.add(
         method="GET",
-        url="https://zsapi.zscaler.net/api/v1/urlCategories?customOnly=False",
+        url="https://zsapi.zscaler.net/api/v1/urlCategories?customOnly=False&includeOnlyUrlKeywordCounts=False",
         json=url_categories,
         status=200,
     )
@@ -190,7 +228,6 @@ def test_update_url_category(zia, custom_categories):
 
 @responses.activate
 def test_add_urls_to_category(zia, custom_categories):
-
     responses.add(
         method="GET",
         url="https://zsapi.zscaler.net/api/v1/urlCategories/CUSTOM_02",