ismijnipverweg/app/main.py

'''
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))