cookbooks/planet/templates/default/apache-s3-ip2region.erb

   1 #!/usr/bin/env python3
   2
   3 import sys
   4 import os
   5 import geoip2.database
   6 import ipaddress
   7
   8 # Constants
   9 DB_PATH = "<%= node[:geoipupdate][:directory] %>/GeoLite2-Country.mmdb"
  10
  11 # Default region when continent doesn't match any in the dictionary
  12 DEFAULT_REGION = "eu-central-1"
  13
  14 # Mapping of continents to AWS regions
  15 CONTINENT_TO_AWS_REGION = {
  16     "NA": "us-west-2", # North America
  17     "OC": "us-west-2", # Oceania
  18     "SA": "us-west-2", # South America
  19 }
  20
  21 # Global to store last known modification time and database reader
  22 last_mod_time = None
  23 reader = None
  24
  25 def is_valid_ip(ip_str):
  26     """Check if a string is a valid IPv4 or IPv6 address."""
  27     try:
  28         ipaddress.ip_address(ip_str)
  29         return True
  30     except ValueError:
  31         return False
  32
  33 def get_reader():
  34     """Get the geoip2 database reader. Reload if the DB file has changed."""
  35     global last_mod_time
  36     global reader
  37
  38     if not os.path.exists(DB_PATH):
  39       return None  # Database file missing
  40
  41     current_mod_time = os.path.getmtime(DB_PATH)
  42
  43     # If file has changed or reader isn't initialized, reload it
  44     if reader is None or current_mod_time != last_mod_time:
  45         if reader:
  46             reader.close()  # Close the existing reader before reinitializing
  47         reader = geoip2.database.Reader(DB_PATH)
  48         last_mod_time = current_mod_time
  49
  50     return reader
  51
  52 def get_continent_from_ip(ip_address):
  53     """Return the continent for a given IP address."""
  54     if not is_valid_ip(ip_address):
  55         return None
  56     reader = get_reader()
  57     if reader is None:
  58       return None  # No continent as DB is missing
  59     try:
  60         response = reader.country(ip_address)
  61         return response.continent.code
  62     except:
  63         return None  # Indicates invalid IP address or other issues
  64
  65 def determine_aws_region(continent_code):
  66     """Determine AWS region based on the continent code using a dictionary."""
  67     return CONTINENT_TO_AWS_REGION.get(continent_code, DEFAULT_REGION)
  68
  69 def main():
  70     """Main function to process IP addresses from stdin and return AWS regions."""
  71     for line in sys.stdin:
  72         ip_address = line.strip()
  73
  74         continent_code = get_continent_from_ip(ip_address)
  75         aws_region = determine_aws_region(continent_code)
  76
  77         sys.stdout.write(f"{aws_region}\n")
  78         sys.stdout.flush()
  79
  80 if __name__ == "__main__":
  81     main()