base.py 8.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227
  1. import socket
  2. import geoip2.database
  3. from django.conf import settings
  4. from django.core.exceptions import ValidationError
  5. from django.core.validators import validate_ipv46_address
  6. from django.utils._os import to_path
  7. from .resources import City, Country
  8. # Creating the settings dictionary with any settings, if needed.
  9. GEOIP_SETTINGS = {
  10. 'GEOIP_PATH': getattr(settings, 'GEOIP_PATH', None),
  11. 'GEOIP_CITY': getattr(settings, 'GEOIP_CITY', 'GeoLite2-City.mmdb'),
  12. 'GEOIP_COUNTRY': getattr(settings, 'GEOIP_COUNTRY', 'GeoLite2-Country.mmdb'),
  13. }
  14. class GeoIP2Exception(Exception):
  15. pass
  16. class GeoIP2:
  17. # The flags for GeoIP memory caching.
  18. # Try MODE_MMAP_EXT, MODE_MMAP, MODE_FILE in that order.
  19. MODE_AUTO = 0
  20. # Use the C extension with memory map.
  21. MODE_MMAP_EXT = 1
  22. # Read from memory map. Pure Python.
  23. MODE_MMAP = 2
  24. # Read database as standard file. Pure Python.
  25. MODE_FILE = 4
  26. # Load database into memory. Pure Python.
  27. MODE_MEMORY = 8
  28. cache_options = frozenset((MODE_AUTO, MODE_MMAP_EXT, MODE_MMAP, MODE_FILE, MODE_MEMORY))
  29. # Paths to the city & country binary databases.
  30. _city_file = ''
  31. _country_file = ''
  32. # Initially, pointers to GeoIP file references are NULL.
  33. _city = None
  34. _country = None
  35. def __init__(self, path=None, cache=0, country=None, city=None):
  36. """
  37. Initialize the GeoIP object. No parameters are required to use default
  38. settings. Keyword arguments may be passed in to customize the locations
  39. of the GeoIP datasets.
  40. * path: Base directory to where GeoIP data is located or the full path
  41. to where the city or country data files (*.mmdb) are located.
  42. Assumes that both the city and country data sets are located in
  43. this directory; overrides the GEOIP_PATH setting.
  44. * cache: The cache settings when opening up the GeoIP datasets. May be
  45. an integer in (0, 1, 2, 4, 8) corresponding to the MODE_AUTO,
  46. MODE_MMAP_EXT, MODE_MMAP, MODE_FILE, and MODE_MEMORY,
  47. `GeoIPOptions` C API settings, respectively. Defaults to 0,
  48. meaning MODE_AUTO.
  49. * country: The name of the GeoIP country data file. Defaults to
  50. 'GeoLite2-Country.mmdb'; overrides the GEOIP_COUNTRY setting.
  51. * city: The name of the GeoIP city data file. Defaults to
  52. 'GeoLite2-City.mmdb'; overrides the GEOIP_CITY setting.
  53. """
  54. # Checking the given cache option.
  55. if cache in self.cache_options:
  56. self._cache = cache
  57. else:
  58. raise GeoIP2Exception('Invalid GeoIP caching option: %s' % cache)
  59. # Getting the GeoIP data path.
  60. path = path or GEOIP_SETTINGS['GEOIP_PATH']
  61. if not path:
  62. raise GeoIP2Exception('GeoIP path must be provided via parameter or the GEOIP_PATH setting.')
  63. path = to_path(path)
  64. if path.is_dir():
  65. # Constructing the GeoIP database filenames using the settings
  66. # dictionary. If the database files for the GeoLite country
  67. # and/or city datasets exist, then try to open them.
  68. country_db = path / (country or GEOIP_SETTINGS['GEOIP_COUNTRY'])
  69. if country_db.is_file():
  70. self._country = geoip2.database.Reader(str(country_db), mode=cache)
  71. self._country_file = country_db
  72. city_db = path / (city or GEOIP_SETTINGS['GEOIP_CITY'])
  73. if city_db.is_file():
  74. self._city = geoip2.database.Reader(str(city_db), mode=cache)
  75. self._city_file = city_db
  76. if not self._reader:
  77. raise GeoIP2Exception('Could not load a database from %s.' % path)
  78. elif path.is_file():
  79. # Otherwise, some detective work will be needed to figure out
  80. # whether the given database path is for the GeoIP country or city
  81. # databases.
  82. reader = geoip2.database.Reader(str(path), mode=cache)
  83. db_type = reader.metadata().database_type
  84. if db_type.endswith('City'):
  85. # GeoLite City database detected.
  86. self._city = reader
  87. self._city_file = path
  88. elif db_type.endswith('Country'):
  89. # GeoIP Country database detected.
  90. self._country = reader
  91. self._country_file = path
  92. else:
  93. raise GeoIP2Exception('Unable to recognize database edition: %s' % db_type)
  94. else:
  95. raise GeoIP2Exception('GeoIP path must be a valid file or directory.')
  96. @property
  97. def _reader(self):
  98. return self._country or self._city
  99. @property
  100. def _country_or_city(self):
  101. if self._country:
  102. return self._country.country
  103. else:
  104. return self._city.city
  105. def __del__(self):
  106. # Cleanup any GeoIP file handles lying around.
  107. if self._reader:
  108. self._reader.close()
  109. def __repr__(self):
  110. meta = self._reader.metadata()
  111. version = '[v%s.%s]' % (meta.binary_format_major_version, meta.binary_format_minor_version)
  112. return '<%(cls)s %(version)s _country_file="%(country)s", _city_file="%(city)s">' % {
  113. 'cls': self.__class__.__name__,
  114. 'version': version,
  115. 'country': self._country_file,
  116. 'city': self._city_file,
  117. }
  118. def _check_query(self, query, country=False, city=False, city_or_country=False):
  119. "Check the query and database availability."
  120. # Making sure a string was passed in for the query.
  121. if not isinstance(query, str):
  122. raise TypeError('GeoIP query must be a string, not type %s' % type(query).__name__)
  123. # Extra checks for the existence of country and city databases.
  124. if city_or_country and not (self._country or self._city):
  125. raise GeoIP2Exception('Invalid GeoIP country and city data files.')
  126. elif country and not self._country:
  127. raise GeoIP2Exception('Invalid GeoIP country data file: %s' % self._country_file)
  128. elif city and not self._city:
  129. raise GeoIP2Exception('Invalid GeoIP city data file: %s' % self._city_file)
  130. # Return the query string back to the caller. GeoIP2 only takes IP addresses.
  131. try:
  132. validate_ipv46_address(query)
  133. except ValidationError:
  134. query = socket.gethostbyname(query)
  135. return query
  136. def city(self, query):
  137. """
  138. Return a dictionary of city information for the given IP address or
  139. Fully Qualified Domain Name (FQDN). Some information in the dictionary
  140. may be undefined (None).
  141. """
  142. enc_query = self._check_query(query, city=True)
  143. return City(self._city.city(enc_query))
  144. def country_code(self, query):
  145. "Return the country code for the given IP Address or FQDN."
  146. enc_query = self._check_query(query, city_or_country=True)
  147. return self.country(enc_query)['country_code']
  148. def country_name(self, query):
  149. "Return the country name for the given IP Address or FQDN."
  150. enc_query = self._check_query(query, city_or_country=True)
  151. return self.country(enc_query)['country_name']
  152. def country(self, query):
  153. """
  154. Return a dictionary with the country code and name when given an
  155. IP address or a Fully Qualified Domain Name (FQDN). For example, both
  156. '24.124.1.80' and 'djangoproject.com' are valid parameters.
  157. """
  158. # Returning the country code and name
  159. enc_query = self._check_query(query, city_or_country=True)
  160. return Country(self._country_or_city(enc_query))
  161. # #### Coordinate retrieval routines ####
  162. def coords(self, query, ordering=('longitude', 'latitude')):
  163. cdict = self.city(query)
  164. if cdict is None:
  165. return None
  166. else:
  167. return tuple(cdict[o] for o in ordering)
  168. def lon_lat(self, query):
  169. "Return a tuple of the (longitude, latitude) for the given query."
  170. return self.coords(query)
  171. def lat_lon(self, query):
  172. "Return a tuple of the (latitude, longitude) for the given query."
  173. return self.coords(query, ('latitude', 'longitude'))
  174. def geos(self, query):
  175. "Return a GEOS Point object for the given query."
  176. ll = self.lon_lat(query)
  177. if ll:
  178. from django.contrib.gis.geos import Point
  179. return Point(ll, srid=4326)
  180. else:
  181. return None
  182. # #### GeoIP Database Information Routines ####
  183. @property
  184. def info(self):
  185. "Return information about the GeoIP library and databases in use."
  186. meta = self._reader.metadata()
  187. return 'GeoIP Library:\n\t%s.%s\n' % (meta.binary_format_major_version, meta.binary_format_minor_version)
  188. @classmethod
  189. def open(cls, full_path, cache):
  190. return GeoIP2(full_path, cache)