From 734ac03c650473a82400d14ccacf4a53a3820fa1 Mon Sep 17 00:00:00 2001
From: Ruben van Staveren <ruben@verweg.com>
Date: Tue, 12 Dec 2023 14:46:14 +0100
Subject: [PATCH] Add post method and let the get method reuse that code

---
 app/main.py | 50 +++++++++++++++++++++++++++++++-------------------
 1 file changed, 31 insertions(+), 19 deletions(-)

diff --git a/app/main.py b/app/main.py
index 93b6a82..7e007dd 100644
--- a/app/main.py
+++ b/app/main.py
@@ -5,7 +5,7 @@ from ipaddress import IPv4Address, IPv4Network, IPv6Address, IPv6Network
 from typing import Annotated, Optional, Union
 
 import geoip2.database
-from fastapi import FastAPI, Path
+from fastapi import FastAPI, Path, Body
 from pydantic import BaseModel
 
 app = FastAPI()
@@ -13,6 +13,9 @@ 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):
+    ip: Union[IPv6Address,IPv4Address]
+
 class Locality(BaseModel):
     '''
     Locality data
@@ -32,27 +35,36 @@ class GeoLocation(BaseModel):
     network: Union[IPv6Network,IPv4Network,None]
     locality: Locality
 
+@app.post("/")
+async def root_post(ipaddresses: Annotated[list[IPAddressParam],
+                                    Body(title="The IPAddresses to geolocate")]
+               ) -> list[GeoLocation]:
+    geolocations = []
+    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
+                )
+            ))
+    return geolocations
+
+
 @app.get("/{ipaddress}")
-async def root(ipaddress: Annotated[Union[IPv4Address,IPv6Address],
+async def root_get(ipaddress: Annotated[Union[IPv4Address,IPv6Address],
                                     Path(title="The IPAddress to geolocate")]
                ) -> GeoLocation:
     '''
     Look up geolocation for ip in path parameter
     '''
-    with (geoip2.database.Reader(GEOLITE2_ASN_DB) as reader_asn,
-            geoip2.database.Reader(GEOLITE2_CITY_DB) as reader_city):
-        asn_data = reader_asn.asn(ipaddress)
-        city_data = reader_city.city(ipaddress)
-
-        return GeoLocation(
-            ip=ipaddress,
-            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
-            )
-        )
+    return (await root_post([IPAddressParam(ip=ipaddress)])).pop()