diff --git a/app/main.py b/app/main.py index 8feb465..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,20 +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")] - ) -> list[GeoLocation]: +async def root_post(ipaddresses: Annotated[ + list[IPAddressParam], + Body(title="The IPAddresses to geolocate")], + response: Response + ) -> list[GeoLocation]: ''' Return GeoLocation item(s) for a list of IPAddressParam objects ''' @@ -72,26 +78,31 @@ async def root_post(ipaddresses: Annotated[list[IPAddressParam], ip=ipaddress.ip ) ) + if geolocations: + response.headers['Cache-Control'] = 'private, max-age=604800' return geolocations @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 ''' - 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() 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))