#!/bin/bash
set -euo pipefail
LICENSE_KEY=${LICENSE_KEY:?LICENSE_KEY missing}
TMPDIR=$(mktemp -p /tmp -d haproxy_country.XXX)

curl --silent \
  --output "$TMPDIR/geoip.zip" \
  "https://download.maxmind.com/app/geoip_download?edition_id=GeoLite2-Country-CSV&license_key=${LICENSE_KEY}&suffix=zip"

unzip -j "$TMPDIR/geoip.zip" -d "$TMPDIR" -x '*.txt'

cd "$TMPDIR"

# Create an array of country codes using the first column of
# GeoLite2-Country-Locations-en.csv as the indices and the fifth column as the
# values
# Use sed to skip the first line
declare -A country_codes
while IFS=',' read -r geoname_id _ _ _ country_iso_code _ _; do
  country_codes[$geoname_id]=$country_iso_code
done < <(sed '1d' GeoLite2-Country-Locations-en.csv)

# Process the blocks file, replacing country identifiers with the corresponding
# country codes
# Use sed to skip the first line
while IFS=',' read -r network geoname_id registered_country_geoname_id _ _ _; do
  # If geoname_id is not present, use registered_country_geoname_id as a substitute
  # Or if registered_country_geoname_id is not present, use whois
  if [[ -z $geoname_id ]]; then
    if [[ -n $registered_country_geoname_id ]]; then
      geoname_id=$registered_country_geoname_id
    else
      country_code=$(whois -h whois.cymru.com "-v $network" | tail -n1 | awk -F'|' '{print $4}' | tr -d ' ')
      # Convert country code to GeoLite country code
      geo_country_code=$(grep "$country_code" GeoLite2-Country-Locations-en.csv | awk -F',' '{print $1}')
      geoname_id=$geo_country_code
    fi
  fi
  echo "$network" >> "${country_codes[$geoname_id]}.txt"
done < <(sed '1d' GeoLite2-Country-Blocks-IPv4.csv)

while IFS=',' read -r network geoname_id registered_country_geoname_id _ _ _; do
  # If geoname_id is not present, use registered_country_geoname_id as a substitute
  # Or if registered_country_geoname_id is not present, use whois
  if [[ -z $geoname_id ]]; then
    if [[ -n $registered_country_geoname_id ]]; then
      geoname_id=$registered_country_geoname_id
    else
      country_code=$(whois -h whois.cymru.com "-v $network" | tail -n1 | awk -F'|' '{print $4}' | tr -d ' ')
      # Convert country code to GeoLite country code
      geo_country_code=$(grep "$country_code" GeoLite2-Country-Locations-en.csv | awk -F',' '{print $1}')
      geoname_id=$geo_country_code
    fi
  fi
  echo "$network" >> "${country_codes[$geoname_id]}.txt"
done < <(sed '1d' GeoLite2-Country-Blocks-IPv6.csv)

rm -f /etc/haproxy/country/*.txt
cp ./*.txt /etc/haproxy/country/

systemctl reload haproxy

cd - > /dev/null
rm -rf "$TMPDIR"