From 0b4322d671838f7659957a04a1ebc9c24cc7eb2b Mon Sep 17 00:00:00 2001
From: Ruben van Staveren <ruben@verweg.com>
Date: Mon, 18 Dec 2023 14:59:27 +0100
Subject: [PATCH] Handle AddressNotFoundError and allow empty placeholder
 GeoLocation for failed queries

---
 app/main.py | 63 ++++++++++++++++++++++++++++++++---------------------
 1 file changed, 38 insertions(+), 25 deletions(-)

diff --git a/app/main.py b/app/main.py
index d7651ee..6322685 100644
--- a/app/main.py
+++ b/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()