Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 29 additions & 32 deletions ipinfo/handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,33 +2,34 @@
Main API client handler for fetching data from the IPinfo service.
"""

from ipaddress import IPv4Address, IPv6Address
import time
from ipaddress import IPv4Address, IPv6Address

import requests

from .error import APIError
from . import handler_utils
from .bogon import is_bogon
from .cache.default import DefaultCache
from .data import (
continents,
countries,
countries_currencies,
countries_flags,
eu_countries,
)
from .details import Details
from .error import APIError
from .exceptions import RequestQuotaExceededError, TimeoutExceededError
from .handler_utils import (
API_URL,
RESPROXY_API_URL,
BATCH_MAX_SIZE,
BATCH_REQ_TIMEOUT_DEFAULT,
CACHE_MAXSIZE,
CACHE_TTL,
REQUEST_TIMEOUT_DEFAULT,
BATCH_REQ_TIMEOUT_DEFAULT,
RESPROXY_API_URL,
cache_key,
)
from . import handler_utils
from .bogon import is_bogon
from .data import (
continents,
countries,
countries_currencies,
eu_countries,
countries_flags,
is_prefixed_lookup,
)


Expand Down Expand Up @@ -91,9 +92,7 @@ def getDetails(self, ip_address=None, timeout=None):
# If the supplied IP address uses the objects defined in the built-in
# module ipaddress extract the appropriate string notation before
# formatting the URL.
if isinstance(ip_address, IPv4Address) or isinstance(
ip_address, IPv6Address
):
if isinstance(ip_address, IPv4Address) or isinstance(ip_address, IPv6Address):
ip_address = ip_address.exploded

# check if bogon.
Expand Down Expand Up @@ -125,11 +124,11 @@ def getDetails(self, ip_address=None, timeout=None):
raise RequestQuotaExceededError()
if response.status_code >= 400:
error_code = response.status_code
content_type = response.headers.get('Content-Type')
if content_type == 'application/json':
content_type = response.headers.get("Content-Type")
if content_type == "application/json":
error_response = response.json()
else:
error_response = {'error': response.text}
error_response = {"error": response.text}
raise APIError(error_code, error_response)
details = response.json()

Expand Down Expand Up @@ -196,7 +195,6 @@ def getResproxy(self, ip_address, timeout=None):

return Details(details)


def getBatchDetails(
self,
ip_addresses,
Expand Down Expand Up @@ -251,7 +249,11 @@ def getBatchDetails(
):
ip_address = ip_address.exploded

if ip_address and is_bogon(ip_address):
if (
ip_address
and not is_prefixed_lookup(ip_address)
and is_bogon(ip_address)
):
details = {}
details["ip"] = ip_address
details["bogon"] = True
Expand Down Expand Up @@ -280,10 +282,7 @@ def getBatchDetails(
headers["content-type"] = "application/json"
for i in range(0, len(lookup_addresses), batch_size):
# quit if total timeout is reached.
if (
timeout_total is not None
and time.time() - start_time > timeout_total
):
if timeout_total is not None and time.time() - start_time > timeout_total:
return handler_utils.return_or_fail(
raise_on_fail, TimeoutExceededError(), result
)
Expand All @@ -292,9 +291,7 @@ def getBatchDetails(

# lookup
try:
response = requests.post(
url, json=chunk, headers=headers, **req_opts
)
response = requests.post(url, json=chunk, headers=headers, **req_opts)
except Exception as e:
return handler_utils.return_or_fail(raise_on_fail, e, result)

Expand Down Expand Up @@ -347,9 +344,7 @@ def getMap(self, ips):
url = f"{API_URL}/map?cli=1"
headers = handler_utils.get_headers(None, self.headers)
headers["content-type"] = "application/json"
response = requests.post(
url, json=ip_strs, headers=headers, **req_opts
)
response = requests.post(url, json=ip_strs, headers=headers, **req_opts)
response.raise_for_status()
return response.json()["reportUrl"]

Expand All @@ -370,7 +365,9 @@ def getBatchDetailsIter(
):
ip_address = ip_address.exploded

if ip_address and is_bogon(ip_address):
if is_prefixed_lookup(ip_address):
lookup_addresses.append(ip_address)
elif ip_address and is_bogon(ip_address):
details = {}
details["ip"] = ip_address
details["bogon"] = True
Expand Down
81 changes: 47 additions & 34 deletions ipinfo/handler_async.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,35 +2,36 @@
Main API client asynchronous handler for fetching data from the IPinfo service.
"""

from ipaddress import IPv4Address, IPv6Address
import asyncio
import json
import time
from ipaddress import IPv4Address, IPv6Address

import aiohttp

from .error import APIError
from . import handler_utils
from .bogon import is_bogon
from .cache.default import DefaultCache
from .data import (
continents,
countries,
countries_currencies,
countries_flags,
eu_countries,
)
from .details import Details
from .error import APIError
from .exceptions import RequestQuotaExceededError, TimeoutExceededError
from .handler_utils import (
API_URL,
RESPROXY_API_URL,
BATCH_MAX_SIZE,
BATCH_REQ_TIMEOUT_DEFAULT,
CACHE_MAXSIZE,
CACHE_TTL,
REQUEST_TIMEOUT_DEFAULT,
BATCH_REQ_TIMEOUT_DEFAULT,
RESPROXY_API_URL,
cache_key,
)
from . import handler_utils
from .bogon import is_bogon
from .data import (
continents,
countries,
countries_currencies,
eu_countries,
countries_flags,
is_prefixed_lookup,
)


Expand Down Expand Up @@ -117,9 +118,7 @@ async def getDetails(self, ip_address=None, timeout=None):
# If the supplied IP address uses the objects defined in the built-in
# module ipaddress, extract the appropriate string notation before
# formatting the URL.
if isinstance(ip_address, IPv4Address) or isinstance(
ip_address, IPv6Address
):
if isinstance(ip_address, IPv4Address) or isinstance(ip_address, IPv6Address):
ip_address = ip_address.exploded

# check if bogon.
Expand Down Expand Up @@ -147,11 +146,11 @@ async def getDetails(self, ip_address=None, timeout=None):
raise RequestQuotaExceededError()
if resp.status >= 400:
error_code = resp.status
content_type = resp.headers.get('Content-Type')
if content_type == 'application/json':
content_type = resp.headers.get("Content-Type")
if content_type == "application/json":
error_response = await resp.json()
else:
error_response = {'error': resp.text()}
error_response = {"error": resp.text()}
raise APIError(error_code, error_response)
details = await resp.json()

Expand Down Expand Up @@ -277,11 +276,19 @@ async def getBatchDetails(
):
ip_address = ip_address.exploded

try:
cached_ipaddr = self.cache[cache_key(ip_address)]
result[ip_address] = cached_ipaddr
except KeyError:
lookup_addresses.append(ip_address)
if (
ip_address
and not is_prefixed_lookup(ip_address)
and is_bogon(ip_address)
):
details = {"ip": ip_address, "bogon": True}
result[ip_address] = Details(details)
else:
try:
cached_ipaddr = self.cache[cache_key(ip_address)]
result[ip_address] = cached_ipaddr
except KeyError:
lookup_addresses.append(ip_address)

# all in cache - return early.
if not lookup_addresses:
Expand All @@ -296,22 +303,24 @@ async def getBatchDetails(
headers = handler_utils.get_headers(self.access_token, self.headers)
headers["content-type"] = "application/json"

# prepare coroutines that will make reqs and update results.
# prepare tasks that will make reqs and update results.
reqs = [
self._do_batch_req(
lookup_addresses[i : i + batch_size],
url,
headers,
timeout_per_batch,
raise_on_fail,
result,
asyncio.ensure_future(
self._do_batch_req(
lookup_addresses[i : i + batch_size],
url,
headers,
timeout_per_batch,
raise_on_fail,
result,
)
)
for i in range(0, len(lookup_addresses), batch_size)
]

try:
_, pending = await asyncio.wait(
{*reqs},
reqs,
timeout=timeout_total,
return_when=asyncio.FIRST_EXCEPTION,
)
Expand Down Expand Up @@ -404,7 +413,11 @@ async def getBatchDetailsIter(
):
ip_address = ip_address.exploded

if ip_address and is_bogon(ip_address):
if (
ip_address
and not is_prefixed_lookup(ip_address)
and is_bogon(ip_address)
):
details = {"ip": ip_address, "bogon": True}
yield Details(details)
else:
Expand Down
10 changes: 10 additions & 0 deletions ipinfo/handler_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -136,3 +136,13 @@ def cache_key(k):
Transforms a user-input key into a versioned cache key.
"""
return f"{k}:{CACHE_KEY_VSN}"


def is_prefixed_lookup(ip_address):
"""
Check if the address is a prefixed batch lookup (e.g., "resproxy/1.2.3.4",
"lookup/8.8.8.8", "domains/google.com").

Prefixed lookups skip bogon checking as they are not plain IP addresses.
"""
return isinstance(ip_address, str) and "/" in ip_address
59 changes: 58 additions & 1 deletion tests/handler_async_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,9 @@
import sys

import aiohttp
import ipinfo
import pytest

import ipinfo
from ipinfo import handler_utils
from ipinfo.cache.default import DefaultCache
from ipinfo.details import Details
Expand Down Expand Up @@ -367,3 +368,59 @@ def mock_get(*args, **kwargs):
# Verify only one API call was made (second was cached)
assert call_count == 1
await handler.deinit()


class MockBatchResponse(MockResponse):
"""MockResponse with raise_for_status for batch endpoint mocking."""

def raise_for_status(self):
if self.status >= 400:
raise Exception(f"HTTP {self.status}")


@pytest.mark.asyncio
async def test_get_batch_details_with_resproxy(monkeypatch):
"""Prefixed lookups like 'resproxy/IP' should not crash in async getBatchDetails."""
mock_api_response = {
"resproxy/1.2.3.4": {"ip": "1.2.3.4", "service": "example"},
"8.8.8.8": {"ip": "8.8.8.8", "country": "US"},
}

async def mock_post(*args, **kwargs):
return MockBatchResponse(
json.dumps(mock_api_response),
200,
{"Content-Type": "application/json"},
)

handler = AsyncHandler("test_token")
handler._ensure_aiohttp_ready()
monkeypatch.setattr(handler.httpsess, "post", mock_post)
result = await handler.getBatchDetails(["resproxy/1.2.3.4", "8.8.8.8"])
assert "resproxy/1.2.3.4" in result
assert "8.8.8.8" in result
await handler.deinit()


@pytest.mark.asyncio
async def test_get_batch_details_mixed_resproxy_and_bogon(monkeypatch):
"""Async getBatchDetails: mixing prefixed, plain, and bogon IPs."""
mock_api_response = {
"resproxy/1.2.3.4": {"ip": "1.2.3.4", "service": "ex"},
"8.8.8.8": {"ip": "8.8.8.8", "country": "US"},
}

async def mock_post(*args, **kwargs):
return MockBatchResponse(
json.dumps(mock_api_response),
200,
{"Content-Type": "application/json"},
)

handler = AsyncHandler("test_token")
handler._ensure_aiohttp_ready()
monkeypatch.setattr(handler.httpsess, "post", mock_post)
result = await handler.getBatchDetails(["resproxy/1.2.3.4", "8.8.8.8", "127.0.0.1"])
assert "resproxy/1.2.3.4" in result
assert "8.8.8.8" in result
await handler.deinit()
Loading