Merge pull request #82 from ipinfo/iterative-batch-ips

Adding iterative batch function for IPs

Author: UmanShahzad

ipinfo/handler.py

@@ -201,10 +201,10 @@ def getBatchDetails(
             batch_size = BATCH_MAX_SIZE
 
         result = {}
+        lookup_addresses = []
 
         # pre-populate with anything we've got in the cache, and keep around
         # the IPs not in the cache.
-        lookup_addresses = []
         for ip_address in ip_addresses:
             # if the supplied IP address uses the objects defined in the
             # built-in module ipaddress extract the appropriate string notation
@@ -214,6 +214,14 @@ def getBatchDetails(
             ):
                 ip_address = ip_address.exploded
 
+            if ip_address and is_bogon(ip_address):
+                details = {}
+                details["ip"] = ip_address
+                details["bogon"] = True
+                result[ip_address] = Details(details)
+            else:
+                lookup_addresses.append(ip_address)
+
             try:
                 cached_ipaddr = self.cache[cache_key(ip_address)]
                 result[ip_address] = cached_ipaddr
@@ -309,3 +317,72 @@ def getMap(self, ips):
         )
         response.raise_for_status()
         return response.json()["reportUrl"]
+
+    def getBatchDetailsIter(
+        self,
+        ip_addresses,
+        batch_size=None,
+        raise_on_fail=True,
+    ):
+        if batch_size is None:
+            batch_size = BATCH_MAX_SIZE
+
+        result = {}
+        lookup_addresses = []
+        for ip_address in ip_addresses:
+            if isinstance(ip_address, IPv4Address) or isinstance(
+                ip_address, IPv6Address
+            ):
+                ip_address = ip_address.exploded
+
+            if ip_address and is_bogon(ip_address):
+                details = {}
+                details["ip"] = ip_address
+                details["bogon"] = True
+                yield Details(details)
+            else:
+                lookup_addresses.append(ip_address)
+
+            try:
+                cached_ipaddr = self.cache[cache_key(ip_address)]
+                result[ip_address] = cached_ipaddr
+            except KeyError:
+                lookup_addresses.append(ip_address)
+
+        # all in cache - exit early.
+        if len(lookup_addresses) == 0:
+            raise StopIteration(result.items())
+
+        url = API_URL + "/batch"
+        headers = handler_utils.get_headers(self.access_token, self.headers)
+        headers["content-type"] = "application/json"
+        for i in range(0, len(lookup_addresses), batch_size):
+            batch = lookup_addresses[i : i + batch_size]
+
+            try:
+                response = requests.post(url, json=batch, headers=headers)
+            except Exception as e:
+                raise e
+
+            try:
+                if response.status_code == 429:
+                    raise RequestQuotaExceededError()
+                response.raise_for_status()
+            except Exception as e:
+                return handler_utils.return_or_fail(raise_on_fail, e)
+
+            details = response.json()
+
+            # format & cache
+            handler_utils.format_details(
+                details,
+                self.countries,
+                self.eu_countries,
+                self.countries_flags,
+                self.countries_currencies,
+                self.continents,
+            )
+            for ip in batch:
+                detail = details.get(ip)
+                self.cache[cache_key(ip)] = detail
+                yield detail

ipinfo/handler_async.py

@@ -351,3 +351,69 @@ def _ensure_aiohttp_ready(self):
 
         timeout = aiohttp.ClientTimeout(total=self.request_options["timeout"])
         self.httpsess = aiohttp.ClientSession(timeout=timeout)
+
+    async def getBatchDetailsIter(
+        self,
+        ip_addresses,
+        batch_size=None,
+        raise_on_fail=True,
+    ):
+        if batch_size is None:
+            batch_size = BATCH_MAX_SIZE
+
+        results = {}
+        lookup_addresses = []
+        for ip_address in ip_addresses:
+            if isinstance(ip_address, IPv4Address) or isinstance(
+                ip_address, IPv6Address
+            ):
+                ip_address = ip_address.exploded
+
+            if ip_address and is_bogon(ip_address):
+                details = {}
+                details["ip"] = ip_address
+                details["bogon"] = True
+                yield Details(details)
+            else:
+                lookup_addresses.append(ip_address)
+
+            try:
+                cached_ipaddr = self.cache[cache_key(ip_address)]
+                results[ip_address] = cached_ipaddr
+            except KeyError:
+                lookup_addresses.append(ip_address)
+
+        if len(lookup_addresses) == 0:
+            yield results.items()
+
+        url = API_URL + "/batch"
+        headers = handler_utils.get_headers(self.access_token, self.headers)
+        headers["content-type"] = "application/json"
+
+        async def process_batch(batch):
+            try:
+                async with aiohttp.ClientSession(headers=headers) as session:
+                    response = await session.post(url, json=batch)
+                    response.raise_for_status()
+                    json_response = await response.json()
+                    for ip_address, details in json_response.items():
+                        self.cache[cache_key(ip_address)] = details
+                        results[ip_address] = details
+            except Exception as e:
+                raise e
+
+        for i in range(0, len(lookup_addresses), batch_size):
+            batch = lookup_addresses[i : i + batch_size]
+            await process_batch(batch)
+
+            for ip_address, details in results.items():
+                if isinstance(details, dict):
+                    handler_utils.format_details(
+                        details,
+                        self.countries,
+                        self.eu_countries,
+                        self.countries_flags,
+                        self.countries_currencies,
+                        self.continents,
+                    )
+                yield ip_address, details

ipinfo/handler_utils.py

@@ -80,7 +80,7 @@ def format_details(
     details["country_name"] = countries.get(details.get("country"))
     details["isEU"] = details.get("country") in eu_countries
     details["country_flag_url"] = (
-        COUNTRY_FLAGS_URL + details.get("country") + ".svg"
+        COUNTRY_FLAGS_URL + (details.get("country") or "") + ".svg"
     )
     details["country_flag"] = copy.deepcopy(
         countries_flags.get(details.get("country"))

tests/handler_async_test.py

@@ -145,6 +145,27 @@ async def test_get_batch_details(batch_size):
     await handler.deinit()
 
 
+def _check_iterative_batch_details(ip, details, token):
+    """Helper for iterative batch tests."""
+    assert ip == details.get("ip")
+    assert "country" in details
+    assert "city" in details
+    if token:
+        assert "asn" in details or "anycast" in details
+        assert "company" in details or "org" in details
+        assert "privacy" in details or "anycast" in details
+        assert "abuse" in details or "anycast" in details
+        assert "domains" in details or "anycast" in details
+
+
+@pytest.mark.parametrize("batch_size", [None, 1, 2, 3])
+@pytest.mark.asyncio
+async def test_get_iterative_batch_details(batch_size):
+    handler, token, ips = _prepare_batch_test()
+    async for ips, details in handler.getBatchDetailsIter(ips, batch_size):
+        _check_iterative_batch_details(ips, details, token)
+
+
 @pytest.mark.parametrize("batch_size", [None, 1, 2, 3])
 @pytest.mark.asyncio
 async def test_get_batch_details_total_timeout(batch_size):

tests/handler_test.py

@@ -131,6 +131,19 @@ def _check_batch_details(ips, details, token):
             assert "domains" in d
 
 
+def _check_iterative_batch_details(details, token):
+    """Helper for iterative batch tests."""
+    assert "ip" in details, "Key 'ip' not found in details"
+    assert "country" in details, "Key 'country' not found in details"
+    assert "city" in details, "Key 'city' not found in details"
+    if token:
+        assert "asn" in details, "Key 'asn' not found in details"
+        assert "company" in details, "Key 'company' not found in details"
+        assert "privacy" in details, "Key 'privacy' not found in details"
+        assert "abuse" in details, "Key 'abuse' not found in details"
+        assert "domains" in details, "Key 'domains' not found in details"
+
+
 @pytest.mark.parametrize("batch_size", [None, 1, 2, 3])
 def test_get_batch_details(batch_size):
     handler, token, ips = _prepare_batch_test()
@@ -147,6 +160,14 @@ def test_get_batch_details_total_timeout(batch_size):
         )
 
 
+@pytest.mark.parametrize("batch_size", [None, 1, 2, 3])
+def test_get_iterative_batch_details(batch_size):
+    handler, token, ips = _prepare_batch_test()
+    details_iterator = handler.getBatchDetailsIter(ips, batch_size=batch_size)
+    for details in details_iterator:
+        _check_iterative_batch_details(details, token)
+
+
 #############
 # MAP TESTS
 #############
@@ -169,3 +190,10 @@ def test_bogon_details():
     details = handler.getDetails("127.0.0.1")
     assert isinstance(details, Details)
     assert details.all == {"bogon": True, "ip": "127.0.0.1"}
+
+
+def test_iterative_bogon_details():
+    token = os.environ.get("IPINFO_TOKEN", "")
+    handler = Handler(token)
+    details = next(handler.getBatchDetailsIter(["127.0.0.1"]))
+    assert details.all == {"bogon": True, "ip": "127.0.0.1"}