]> git.openstreetmap.org Git - chef.git/blob - cookbooks/planet/templates/default/apache-s3-ip2region.erb
Add redirects to planet S3
[chef.git] / 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()