From d090e90ee20c1a53028c92963b496d865e498b19 Mon Sep 17 00:00:00 2001
From: Ruben van Staveren <ruben@verweg.com>
Date: Mon, 6 May 2024 17:24:27 +0200
Subject: [PATCH 1/3] Add Cache-Control headers

---
 app/main.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/app/main.py b/app/main.py
index 8feb465..16e8a55 100644
--- a/app/main.py
+++ b/app/main.py
@@ -42,7 +42,8 @@ class GeoLocation(BaseModel):
 
 @app.post("/")
 async def root_post(ipaddresses: Annotated[list[IPAddressParam],
-                                    Body(title="The IPAddresses to geolocate")]
+                                    Body(title="The IPAddresses to geolocate")],
+                    response: Response
                ) -> list[GeoLocation]:
     '''
     Return GeoLocation item(s) for a list of IPAddressParam objects
@@ -72,6 +73,8 @@ async def root_post(ipaddresses: Annotated[list[IPAddressParam],
                     ip=ipaddress.ip
                     )
                 )
+    if geolocations:
+        response.headers['Cache-Control'] = 'private, max-age=604800'
     return geolocations
 
 
@@ -85,6 +88,7 @@ async def root_get(ipaddress: Annotated[Union[IPv4Address,IPv6Address],
     '''
     locations = await root_post([IPAddressParam(ip=ipaddress)])
     if locations:
+        response.headers['Cache-Control'] = 'private, max-age=604800'
         return locations.pop()
     response.status_code = status.HTTP_404_NOT_FOUND
     return GeoLocation()
-- 
2.49.0


From 878a2b6bbf6555dbe4b169adbbec32e8a78ba489 Mon Sep 17 00:00:00 2001
From: Ruben van Staveren <ruben@verweg.com>
Date: Tue, 7 May 2024 16:50:41 +0200
Subject: [PATCH 2/3] Add missing parameter

---
 app/main.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/app/main.py b/app/main.py
index 16e8a55..02dc5a8 100644
--- a/app/main.py
+++ b/app/main.py
@@ -86,7 +86,7 @@ async def root_get(ipaddress: Annotated[Union[IPv4Address,IPv6Address],
     '''
     Look up geolocation for ip in path parameter
     '''
-    locations = await root_post([IPAddressParam(ip=ipaddress)])
+    locations = await root_post([IPAddressParam(ip=ipaddress)], response)
     if locations:
         response.headers['Cache-Control'] = 'private, max-age=604800'
         return locations.pop()
-- 
2.49.0


From 4eb906f6d48a535efdb8d37c7778febab00d4e54 Mon Sep 17 00:00:00 2001
From: Ruben van Staveren <ruben@verweg.com>
Date: Tue, 7 May 2024 17:12:23 +0200
Subject: [PATCH 3/3] Code cleanup

---
 app/main.py | 27 +++++++++++++++++----------
 1 file changed, 17 insertions(+), 10 deletions(-)

diff --git a/app/main.py b/app/main.py
index 02dc5a8..6396921 100644
--- a/app/main.py
+++ b/app/main.py
@@ -15,11 +15,13 @@ app = FastAPI()
 GEOLITE2_ASN_DB = '/usr/local/share/GeoIP/GeoLite2-ASN.mmdb'
 GEOLITE2_CITY_DB = '/usr/local/share/GeoIP/GeoLite2-City.mmdb'
 
+
 class IPAddressParam(BaseModel):
     '''
     Payload entry as used in POST
     '''
-    ip: Union[IPv6Address,IPv4Address]
+    ip: Union[IPv6Address, IPv4Address]
+
 
 class Locality(BaseModel):
     '''
@@ -30,21 +32,24 @@ class Locality(BaseModel):
     continent: Optional[str] = None
     is_eu: bool = False
 
+
 class GeoLocation(BaseModel):
     '''
     Geolocation data model
     '''
-    ip: Optional[Union[IPv6Address,IPv4Address]] = None
+    ip: Optional[Union[IPv6Address, IPv4Address]] = None
     asn: Optional[int] = None
     asn_org: Optional[str] = None
-    network: Optional[Union[IPv6Network,IPv4Network]] = None
+    network: Optional[Union[IPv6Network, IPv4Network]] = None
     locality: Locality = Locality()
 
+
 @app.post("/")
-async def root_post(ipaddresses: Annotated[list[IPAddressParam],
-                                    Body(title="The IPAddresses to geolocate")],
+async def root_post(ipaddresses: Annotated[
+    list[IPAddressParam],
+    Body(title="The IPAddresses to geolocate")],
                     response: Response
-               ) -> list[GeoLocation]:
+                    ) -> list[GeoLocation]:
     '''
     Return GeoLocation item(s) for a list of IPAddressParam objects
     '''
@@ -79,10 +84,11 @@ async def root_post(ipaddresses: Annotated[list[IPAddressParam],
 
 
 @app.get("/{ipaddress}")
-async def root_get(ipaddress: Annotated[Union[IPv4Address,IPv6Address],
-                                    Path(title="The IPAddress to geolocate")],
+async def root_get(ipaddress: Annotated[
+    Union[IPv4Address, IPv6Address],
+    Path(title="The IPAddress to geolocate")],
                    response: Response
-               ) -> GeoLocation:
+                   ) -> GeoLocation:
     '''
     Look up geolocation for ip in path parameter
     '''
@@ -93,9 +99,10 @@ async def root_get(ipaddress: Annotated[Union[IPv4Address,IPv6Address],
     response.status_code = status.HTTP_404_NOT_FOUND
     return GeoLocation()
 
+
 @app.get("/")
 def root_redirect(req: Request) -> RedirectResponse:
     '''
     Redirect empty request using REMOTE_ADDR
     '''
-    return RedirectResponse(url= str(req.url) + str(req.client.host))
+    return RedirectResponse(url=str(req.url) + str(req.client.host))
-- 
2.49.0