From 4c4203a67911cece85fcc335b00101653f134444 Mon Sep 17 00:00:00 2001
From: Ruben van Staveren <ruben@verweg.com>
Date: Tue, 12 Dec 2023 10:43:09 +0100
Subject: [PATCH] Simple FastAPI geolocation app

---
 app/__init__.py |  0
 app/main.py     | 61 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+)
 create mode 100644 app/__init__.py
 create mode 100644 app/main.py

diff --git a/app/__init__.py b/app/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/app/main.py b/app/main.py
new file mode 100644
index 0000000..76908a0
--- /dev/null
+++ b/app/main.py
@@ -0,0 +1,61 @@
+'''
+Simple Geolocation with FastAPI
+'''
+import dataclasses
+from ipaddress import IPv4Address, IPv4Network, IPv6Address, IPv6Network
+from typing import Annotated, Optional, Union
+
+import geoip2.database
+from fastapi import FastAPI, Path
+from pydantic import BaseModel
+
+app = FastAPI()
+
+GEOLITE2_ASN_DB = '/usr/local/share/GeoIP/GeoLite2-ASN.mmdb'
+GEOLITE2_CITY_DB = '/usr/local/share/GeoIP/GeoLite2-City.mmdb'
+
+@dataclasses.dataclass
+class Locality(BaseModel):
+    '''
+    Locality data
+    '''
+    city: Optional[str]
+    country: Optional[str]
+    continent: Optional[str]
+    is_eu: bool
+
+@dataclasses.dataclass
+class GeoLocation(BaseModel):
+    '''
+    Geolocation data model
+    '''
+    ip: Union[IPv6Address,IPv4Address]
+    asn: Optional[int]
+    asn_org: Optional[str]
+    network: Union[IPv6Network,IPv4Network,None]
+    locality: Locality
+
+@app.get("/{ipaddress}")
+async def root(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
+            )
+        )