Add post method and let the get method reuse that code

This commit is contained in:
Ruben van Staveren 2023-12-12 14:46:14 +01:00
parent 49e16b106f
commit 734ac03c65
Signed by: ruben
GPG Key ID: 886F6BECD477A93F

View File

@ -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()