diff --git a/lib/globals.py b/lib/globals.py index a97c4af..452658e 100644 --- a/lib/globals.py +++ b/lib/globals.py @@ -83,6 +83,11 @@ PLAIN_TEXT_AGENTS = [ PLAIN_TEXT_PAGES = [':help', ':bash.function', ':translation', ':iterm2'] +_IPLOCATION_ORDER = os.environ.get( + "WTTR_IPLOCATION_ORDER", + 'geoip,ip2location,ipinfo') +IPLOCATION_ORDER = _IPLOCATION_ORDER.split(',') + _IP2LOCATION_KEY_FILE = os.environ.get( "WTTR_IP2LOCATION_KEY_FILE", os.environ['HOME'] + '/.ip2location.key') @@ -90,6 +95,13 @@ IP2LOCATION_KEY = None if os.path.exists(_IP2LOCATION_KEY_FILE): IP2LOCATION_KEY = open(_IP2LOCATION_KEY_FILE, 'r').read().strip() +_IPINFO_KEY_FILE = os.environ.get( + "WTTR_IPINFO_KEY_FILE", + os.environ['HOME'] + '/.ipinfo.key') +IPINFO_TOKEN = None +if os.path.exists(_IPINFO_KEY_FILE): + IPINFO_TOKEN = open(_IPINFO_KEY_FILE, 'r').read().strip() + _WWO_KEY_FILE = os.environ.get( "WTTR_WWO_KEY_FILE", os.environ['HOME'] + '/.wwo.key') diff --git a/lib/location.py b/lib/location.py index 20f5a12..8388f72 100644 --- a/lib/location.py +++ b/lib/location.py @@ -17,7 +17,7 @@ import requests import geoip2.database from globals import GEOLITE, GEOLOCATOR_SERVICE, IP2LCACHE, IP2LOCATION_KEY, NOT_FOUND_LOCATION, \ - ALIASES, BLACKLIST, IATA_CODES_FILE + ALIASES, BLACKLIST, IATA_CODES_FILE, IPLOCATION_ORDER, IPINFO_TOKEN GEOIP_READER = geoip2.database.Reader(GEOLITE) @@ -88,9 +88,15 @@ def geolocator(location): return None -def ip2location(ip_addr): - "Convert IP address `ip_addr` to a location name" +def ipcachewrite(ip_addr, location): + cached = os.path.join(IP2LCACHE, ip_addr) + if not os.path.exists(IP2LCACHE): + os.makedirs(IP2LCACHE) + with open(cached, 'w') as file: + file.write(location[0] + ';' + location[1]) + +def ipcache(ip_addr): cached = os.path.join(IP2LCACHE, ip_addr) if not os.path.exists(IP2LCACHE): os.makedirs(IP2LCACHE) @@ -98,33 +104,51 @@ def ip2location(ip_addr): location = None if os.path.exists(cached): - location = open(cached, 'r').read() - else: - # if IP2LOCATION_KEY is not set, do not the query, - # because the query wont be processed anyway - if IP2LOCATION_KEY: - try: - ip2location_response = requests\ - .get('http://api.ip2location.com/?ip=%s&key=%s&package=WS3' \ - % (ip_addr, IP2LOCATION_KEY)).text - if ';' in ip2location_response: - open(cached, 'w').write(ip2location_response) - location = ip2location_response - except requests.exceptions.ConnectionError: - pass + location = open(cached, 'r').read().split(';') + return location + +def ip2location(ip_addr): + "Convert IP address `ip_addr` to a location name" + + location = ipcache(ip_addr) + if location: + return location + + # if IP2LOCATION_KEY is not set, do not the query, + # because the query wont be processed anyway + if IP2LOCATION_KEY: + try: + location = requests\ + .get('http://api.ip2location.com/?ip=%s&key=%s&package=WS3' \ + % (ip_addr, IP2LOCATION_KEY)).text + except requests.exceptions.ConnectionError: + pass if location and ';' in location: location = location.split(';')[3], location.split(';')[1] else: location = location, None + if location: + ipcachewrite(ip_addr, location) return location -def get_location(ip_addr): - """ - Return location pair (CITY, COUNTRY) for `ip_addr` - """ +def ipinfo(ip_addr): + location = ipcache(ip_addr) + if location: + return location + if IPINFO_TOKEN: + r = requests.get('https://ipinfo.io/%s/json?token=%s' % + (ip_addr, IPINFO_TOKEN)) + if r.status_code == 200: + location = r.json()["city"], r.json()["country"] + if location: + ipcachewrite(ip_addr, location) + return location + + +def geoip(ip_addr): try: response = GEOIP_READER.city(ip_addr) country = response.country.name @@ -132,7 +156,34 @@ def get_location(ip_addr): except geoip2.errors.AddressNotFoundError: country = None city = None + return city, country +def workaround(city, country): + # workaround for the strange bug with the country name + # maybe some other countries has this problem too + # + # Having these in a separate function will help if this gets to + # be a problem + if country == 'Russian Federation': + country = 'Russia' + return city, country + +def get_location(ip_addr): + """ + Return location pair (CITY, COUNTRY) for `ip_addr` + """ + for method in IPLOCATION_ORDER: + if method == 'geoip': + city, country = geoip(ip_addr) + elif method == 'ip2location': + city, country = ip2location(ip_addr) + elif method == 'ipinfo': + city, country = ipinfo(ip_addr) + else: + print("ERROR: invalid iplocation method speficied: %s" % method) + if city is not None: + city, country = workaround(city, country) + return city, country # # temporary disabled it because of geoip services capcacity # @@ -143,18 +194,9 @@ def get_location(ip_addr): # city = location.raw.get('address', {}).get('city') # except: # pass - if city is None: - city, country = ip2location(ip_addr) - # workaround for the strange bug with the country name - # maybe some other countries has this problem too - if country == 'Russian Federation': - country = 'Russia' - - if city: - return city, country - else: - return NOT_FOUND_LOCATION, None + # No methods resulted in a location - return default + return NOT_FOUND_LOCATION, None def location_canonical_name(location):