Handle AddressNotFoundError and allow empty placeholder GeoLocation for failed queries
This commit is contained in:
parent
949e8582bc
commit
0b4322d671
63
app/main.py
63
app/main.py
@ -5,7 +5,8 @@ from ipaddress import IPv4Address, IPv4Network, IPv6Address, IPv6Network
|
||||
from typing import Annotated, Optional, Union
|
||||
|
||||
import geoip2.database
|
||||
from fastapi import FastAPI, Path, Body
|
||||
from geoip2.errors import AddressNotFoundError
|
||||
from fastapi import FastAPI, Path, Body, Response, status
|
||||
from pydantic import BaseModel
|
||||
|
||||
app = FastAPI()
|
||||
@ -23,20 +24,20 @@ class Locality(BaseModel):
|
||||
'''
|
||||
Locality data
|
||||
'''
|
||||
city: Optional[str]
|
||||
country: Optional[str]
|
||||
continent: Optional[str]
|
||||
is_eu: bool
|
||||
city: Optional[str] = None
|
||||
country: Optional[str] = None
|
||||
continent: Optional[str] = None
|
||||
is_eu: bool = False
|
||||
|
||||
class GeoLocation(BaseModel):
|
||||
'''
|
||||
Geolocation data model
|
||||
'''
|
||||
ip: Union[IPv6Address,IPv4Address]
|
||||
asn: Optional[int]
|
||||
asn_org: Optional[str]
|
||||
network: Union[IPv6Network,IPv4Network,None]
|
||||
locality: Locality
|
||||
ip: Optional[Union[IPv6Address,IPv4Address]] = None
|
||||
asn: Optional[int] = None
|
||||
asn_org: Optional[str] = None
|
||||
network: Optional[Union[IPv6Network,IPv4Network]] = None
|
||||
locality: Locality = Locality()
|
||||
|
||||
@app.post("/")
|
||||
async def root_post(ipaddresses: Annotated[list[IPAddressParam],
|
||||
@ -49,28 +50,40 @@ async def root_post(ipaddresses: Annotated[list[IPAddressParam],
|
||||
with (geoip2.database.Reader(GEOLITE2_ASN_DB) as reader_asn,
|
||||
geoip2.database.Reader(GEOLITE2_CITY_DB) as reader_city):
|
||||
for ipaddress in ipaddresses:
|
||||
asn_data = reader_asn.asn(ipaddress.ip)
|
||||
city_data = reader_city.city(ipaddress.ip)
|
||||
geolocations.append(GeoLocation(
|
||||
ip=ipaddress.ip,
|
||||
asn=asn_data.autonomous_system_number,
|
||||
asn_org=asn_data.autonomous_system_organization,
|
||||
network=asn_data.network,
|
||||
locality=Locality(
|
||||
city=city_data.city.name,
|
||||
country=city_data.country.iso_code,
|
||||
continent=city_data.continent.code,
|
||||
is_eu=city_data.country.is_in_european_union
|
||||
try:
|
||||
asn_data = reader_asn.asn(ipaddress.ip)
|
||||
city_data = reader_city.city(ipaddress.ip)
|
||||
|
||||
geolocations.append(GeoLocation(
|
||||
ip=ipaddress.ip,
|
||||
asn=asn_data.autonomous_system_number,
|
||||
asn_org=asn_data.autonomous_system_organization,
|
||||
network=asn_data.network,
|
||||
locality=Locality(
|
||||
city=city_data.city.name,
|
||||
country=city_data.country.iso_code,
|
||||
continent=city_data.continent.code,
|
||||
is_eu=city_data.country.is_in_european_union
|
||||
)
|
||||
))
|
||||
except AddressNotFoundError:
|
||||
geolocations.append(GeoLocation(
|
||||
ip=ipaddress.ip
|
||||
)
|
||||
)
|
||||
))
|
||||
return geolocations
|
||||
|
||||
|
||||
@app.get("/{ipaddress}")
|
||||
async def root_get(ipaddress: Annotated[Union[IPv4Address,IPv6Address],
|
||||
Path(title="The IPAddress to geolocate")]
|
||||
Path(title="The IPAddress to geolocate")],
|
||||
response: Response
|
||||
) -> GeoLocation:
|
||||
'''
|
||||
Look up geolocation for ip in path parameter
|
||||
'''
|
||||
return (await root_post([IPAddressParam(ip=ipaddress)])).pop()
|
||||
locations = await root_post([IPAddressParam(ip=ipaddress)])
|
||||
if locations:
|
||||
return locations.pop()
|
||||
response.status_code = status.HTTP_404_NOT_FOUND
|
||||
return GeoLocation()
|
||||
|
Loading…
x
Reference in New Issue
Block a user