bridgedb/geo.py

   1 #
   2 #
   3 # This file is part of BridgeDB, a Tor bridge distribution system.
   4 #
   5 # :authors: see AUTHORS file
   6 # :copyright: (c) 2007-2017, The Tor Project, Inc.
   7 # :license: 3-Clause BSD, see LICENSE for licensing information
   8
   9 """
  10 Boilerplate setup for GeoIP. GeoIP allows us to look up the country code
  11 associated with an IP address. This is a "pure" python version which interacts
  12 with the Maxmind GeoIP API (version 1). It requires, in Debian, the libgeoip-dev
  13 and geoip-database packages.
  14 """
  15
  16 import logging
  17 from os.path import isfile
  18
  19 from ipaddr import IPv4Address, IPv6Address
  20
  21 # IPv4 database
  22 GEOIP_DBFILE = '/usr/share/GeoIP/GeoIP.dat'
  23 # IPv6 database
  24 GEOIPv6_DBFILE = '/usr/share/GeoIP/GeoIPv6.dat'
  25 try:
  26     # Make sure we have the database before trying to import the module:
  27     if not (isfile(GEOIP_DBFILE) and isfile(GEOIPv6_DBFILE)):  # pragma: no cover
  28         raise EnvironmentError("Could not find %r. On Debian-based systems, "
  29                                "please install the geoip-database package."
  30                                % GEOIP_DBFILE)
  31
  32     import pygeoip
  33     geoip = pygeoip.GeoIP(GEOIP_DBFILE, flags=pygeoip.MEMORY_CACHE)
  34     geoipv6 = pygeoip.GeoIP(GEOIPv6_DBFILE, flags=pygeoip.MEMORY_CACHE)
  35     logging.info("GeoIP databases loaded")
  36 except Exception as err:  # pragma: no cover
  37     logging.warn("Error while loading geoip module: %r" % err)
  38     geoip = None
  39     geoipv6 = None
  40
  41
  42 def getCountryCode(ip):
  43     """Return the two-letter country code of a given IP address.
  44
  45     :type ip: :class:`ipaddr.IPAddress`
  46     :param ip: An IPv4 OR IPv6 address.
  47     :rtype: ``None`` or str
  48
  49     :returns: If the GeoIP databases are loaded, and the **ip** lookup is
  50         successful, then this returns a two-letter country code.  Otherwise,
  51         this returns ``None``.
  52     """
  53     addr = None
  54     version = None
  55     try:
  56         addr = ip.compressed
  57         version = ip.version
  58     except AttributeError as err:
  59         logging.warn("Wrong type passed to getCountryCode: %s" % str(err))
  60         return None
  61
  62     # GeoIP has two databases: one for IPv4 addresses, and one for IPv6
  63     # addresses. This will ensure we use the correct one.
  64     db = None
  65     # First, make sure we loaded GeoIP properly.
  66     if None in (geoip, geoipv6):
  67         logging.warn("GeoIP databases aren't loaded; can't look up country code")
  68         return None
  69     else:
  70         if version == 4:
  71             db = geoip
  72         else:
  73             db = geoipv6
  74
  75     # Look up the country code of the address.
  76     countryCode = db.country_code_by_addr(addr)
  77     if countryCode:
  78         logging.debug("Looked up country code: %s" % countryCode)
  79         return countryCode
  80     else:
  81         logging.debug("Country code was not detected. IP: %s" % addr)
  82         return None