base.py 8.7 KB

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