diff --git a/docsrc/zs/zia/url_categories.rst b/docsrc/zs/zia/url_categories.rst index d732472..3175c7a 100644 --- a/docsrc/zs/zia/url_categories.rst +++ b/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 diff --git a/pyzscaler/utils.py b/pyzscaler/utils.py index ac86014..459f889 100644 --- a/pyzscaler/utils.py +++ b/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): diff --git a/pyzscaler/zia/url_categories.py b/pyzscaler/zia/url_categories.py index e341408..63ba2e9 100644 --- a/pyzscaler/zia/url_categories.py +++ b/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. diff --git a/tests/zia/test_url_categories.py b/tests/zia/test_url_categories.py index be2f99f..0663ea9 100644 --- a/tests/zia/test_url_categories.py +++ b/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",