src/geoipyupdate/__init__.py

   1 #!/usr/bin/python3
   2 r"""
   3 Simple python replacement for the MaxMind geoipupdate program.
   4 """
   5
   6 # stdlib imports
   7 from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
   8 import gzip
   9 import hashlib
  10 import os
  11 from pathlib import Path
  12 import shutil
  13 import sys
  14 import tempfile
  15 import tomllib
  16
  17 # external imports
  18 import requests
  19
  20
  21 def main():
  22     r"""
  23     The entry point of the geoipyupdate script.
  24     """
  25     # Create an argument parser using our docsctring as its description.
  26     parser = ArgumentParser(description=sys.modules[__name__].__doc__,
  27                             formatter_class=ArgumentDefaultsHelpFormatter)
  28
  29     # XDG_CONFIG_HOME defaults to ~/.config
  30     default_xdgch = str(Path.home() / ".config")
  31     xdgch = os.environ.get("XDG_CONFIG_HOME", default_xdgch)
  32     default_config_file = os.path.join(xdgch,
  33                                        "geoipyupdate",
  34                                        "geoipyupdate.toml")
  35
  36     parser.add_argument('-c',
  37                         '--config-file',
  38                         default=default_config_file,
  39                         help='path to configuration file')
  40
  41     args = parser.parse_args()
  42
  43     # Load/parse the config
  44     with open(args.config_file, "rb") as f:
  45         config = tomllib.load(f)
  46         editions = config["database"]["editions"]
  47         datadir = config["database"]["datadir"]
  48         account_id = config["account"]["account_id"]
  49         license_key = config["account"]["license_key"]
  50
  51     # Impersonate the true client
  52     headers = {'User-Agent': 'geoipupdate/6.1.0'}
  53
  54     # This never changes
  55     url_server = "https://updates.maxmind.com"
  56
  57     for edition in editions:
  58         # The final location of this database, which also happens to
  59         # be where the previous database might be found.
  60         dbfile = os.path.join(datadir, f"{edition}.mmdb")
  61
  62         # Compute the hash of the old database, if there is
  63         # one. Otherwise, leave it blank. This is passed to the server
  64         # who might tell us that the database was 304 Not Modified.
  65         oldhash = ""
  66         if os.path.isfile(dbfile):
  67             with open(dbfile, 'rb') as f:
  68                 oldhash = hashlib.md5(f.read()).hexdigest()
  69
  70         url_path = f"/geoip/databases/{edition}/update?db_md5={oldhash}"
  71         url = f"{url_server}{url_path}"
  72         r = requests.get(url,
  73                          auth=(account_id, license_key),
  74                          headers=headers,
  75                          stream=True,
  76                          timeout=60)
  77
  78         if r.status_code == 304:
  79             # The database hasn't changed since we last downloaded it.
  80             continue
  81
  82         r.raise_for_status()
  83
  84         # Insist on md5 verification of the downloads, i.e. don't handle
  85         # the case where the md5 response header is missing.
  86         xdbmd5 = r.headers["X-Database-MD5"]
  87
  88         # First download the gzipped file to /tmp or wherever. When
  89         # python-3.12 is more widespread, delete_on_close=False might
  90         # be a better alternative, allowing us to use a context
  91         # manager here.
  92         f = tempfile.NamedTemporaryFile(delete=False)
  93         for chunk in r.iter_content(chunk_size=128):
  94             f.write(chunk)
  95         f.close()
  96
  97         # Now gunzip it to a new temporary file (can't simply gunzip in
  98         # place, because now the name would be predictable). We need
  99         # delete=False here because we intend to move this file to the
 100         # datadir.
 101         g = tempfile.NamedTemporaryFile(delete=False)
 102         gdata = gzip.open(f.name, 'rb').read()
 103
 104         # We're done with f. Remove it ASAP in case something goes
 105         # wrong.
 106         os.unlink(f.name)
 107
 108         newhash = hashlib.md5(gdata).hexdigest()
 109         if newhash == xdbmd5:
 110             g.write(gdata)
 111             g.close()
 112         else:
 113             raise ValueError(
 114                 f"{edition} hash doesn't match X-Database-MD5 header"
 115             )
 116
 117         # Overwrite the old database file with the new (gunzipped) one.
 118         shutil.move(g.name, dbfile)