130 lines
3.9 KiB
Python
130 lines
3.9 KiB
Python
'''
|
|
Simple Geolocation with FastAPI
|
|
'''
|
|
import os
|
|
from ipaddress import IPv4Address, IPv4Network, IPv6Address, IPv6Network
|
|
from typing import Annotated, Optional, Union
|
|
|
|
import geoip2.database
|
|
from dotenv import load_dotenv
|
|
from fastapi import Body, FastAPI, Path, Request, Response, status
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from fastapi.responses import RedirectResponse
|
|
from geoip2.errors import AddressNotFoundError
|
|
from pydantic import BaseModel
|
|
|
|
# Load environment variables
|
|
load_dotenv()
|
|
|
|
app = FastAPI()
|
|
|
|
# Configure CORS from environment variables
|
|
cors_origins = os.getenv('CORS_ALLOW_ORIGINS', 'http://localhost')
|
|
allow_origins = [origin.strip() for origin in cors_origins.split(',')
|
|
if origin.strip()]
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=allow_origins,
|
|
allow_credentials=True,
|
|
allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
GEOLITE2_ASN_DB = os.getenv('GEOLITE2_ASN_DB',
|
|
'/usr/local/share/GeoIP/GeoLite2-ASN.mmdb')
|
|
GEOLITE2_CITY_DB = os.getenv('GEOLITE2_CITY_DB',
|
|
'/usr/local/share/GeoIP/GeoLite2-City.mmdb')
|
|
|
|
|
|
class IPAddressParam(BaseModel):
|
|
'''
|
|
Payload entry as used in POST
|
|
'''
|
|
ip: Union[IPv6Address, IPv4Address]
|
|
|
|
|
|
class Locality(BaseModel):
|
|
'''
|
|
Locality data
|
|
'''
|
|
city: Optional[str] = None
|
|
country: Optional[str] = None
|
|
continent: Optional[str] = None
|
|
is_eu: bool = False
|
|
|
|
|
|
class GeoLocation(BaseModel):
|
|
'''
|
|
Geolocation data model
|
|
'''
|
|
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],
|
|
Body(title="The IPAddresses to geolocate")],
|
|
response: Response
|
|
) -> list[GeoLocation]:
|
|
'''
|
|
Return GeoLocation item(s) for a list of IPAddressParam objects
|
|
'''
|
|
geolocations = []
|
|
with (geoip2.database.Reader(GEOLITE2_ASN_DB) as reader_asn,
|
|
geoip2.database.Reader(GEOLITE2_CITY_DB) as reader_city):
|
|
for ipaddress in ipaddresses:
|
|
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
|
|
)
|
|
)
|
|
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")],
|
|
response: Response
|
|
) -> GeoLocation:
|
|
'''
|
|
Look up geolocation for ip in path parameter
|
|
'''
|
|
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))
|