Python: Which country an IP address resides in

which-country.py

from pathlib import Path
from typing import Optional, Union, NewType, cast, Tuple, Iterator
from ipaddress import (
    IPv6Address,
    IPv6Network,
    ip_address,
    ip_network,
    IPv4Address,
    IPv4Network,
)

IPAddress = Union[IPv4Address, IPv6Address]
IPNetwork = Union[IPv4Network, IPv6Network]
CountryCode = NewType("CountryCode", str)


def load_countries(cidr_dir: Path) -> Iterator[Tuple[CountryCode, IPNetwork]]:
    for p in cidr_dir.glob("*.cidr"):
        country_code = cast(CountryCode, p.stem.upper())
        with p.open("r") as fp:
            for line in fp:
                cidr = line.strip()
                yield country_code, ip_network(cidr)


def find_country(ip: IPAddress) -> Optional[CountryCode]:
    sub_dir = f"ipv{ip.version}"
    # database from: https://github.com/herrbischoff/country-ip-blocks
    cidr_dir = Path("country-ip-blocks-master") / sub_dir
    ip_blocks = load_countries(cidr_dir)
    for country, block in ip_blocks:
        if ip in block:
            return country
    else:
        return None


ip = ip_address("1.2.3.4")
country = find_country(ip)
print(country)