src/python/importer.py

   1 #!/usr/bin/python3
   2 ###############################################################################
   3 #                                                                             #
   4 # libloc - A library to determine the location of someone on the Internet     #
   5 #                                                                             #
   6 # Copyright (C) 2020 IPFire Development Team <info@ipfire.org>                #
   7 #                                                                             #
   8 # This library is free software; you can redistribute it and/or               #
   9 # modify it under the terms of the GNU Lesser General Public                  #
  10 # License as published by the Free Software Foundation; either                #
  11 # version 2.1 of the License, or (at your option) any later version.          #
  12 #                                                                             #
  13 # This library is distributed in the hope that it will be useful,             #
  14 # but WITHOUT ANY WARRANTY; without even the implied warranty of              #
  15 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU           #
  16 # Lesser General Public License for more details.                             #
  17 #                                                                             #
  18 ###############################################################################
  19
  20 import gzip
  21 import logging
  22 import urllib.request
  23
  24 # Initialise logging
  25 log = logging.getLogger("location.importer")
  26 log.propagate = 1
  27
  28 class Downloader(object):
  29         def __init__(self):
  30                 self.proxy = None
  31
  32         def set_proxy(self, url):
  33                 """
  34                         Sets a HTTP proxy that is used to perform all requests
  35                 """
  36                 log.info("Using proxy %s" % url)
  37                 self.proxy = url
  38
  39         def request(self, url, data=None):
  40                 req = urllib.request.Request(url, data=data)
  41
  42                 # Configure proxy
  43                 if self.proxy:
  44                         req.set_proxy(self.proxy, "http")
  45
  46                 return DownloaderContext(self, req)
  47
  48
  49 class DownloaderContext(object):
  50         def __init__(self, downloader, request):
  51                 self.downloader = downloader
  52                 self.request = request
  53
  54                 # Save the response object
  55                 self.response = None
  56
  57         def __enter__(self):
  58                 log.info("Retrieving %s..." % self.request.full_url)
  59
  60                 # Send request
  61                 self.response = urllib.request.urlopen(self.request)
  62
  63                 # Log the response headers
  64                 log.debug("Response Headers:")
  65                 for header in self.headers:
  66                         log.debug("     %s: %s" % (header, self.get_header(header)))
  67
  68                 return self
  69
  70         def __exit__(self, type, value, traceback):
  71                 pass
  72
  73         def __iter__(self):
  74                 """
  75                         Makes the object iterable by going through each block
  76                 """
  77                 # Store body
  78                 body = self.body
  79
  80                 while True:
  81                         line = body.readline()
  82                         if not line:
  83                                 break
  84
  85                         # Decode the line
  86                         line = line.decode()
  87
  88                         # Strip the ending
  89                         yield line.rstrip()
  90
  91         @property
  92         def headers(self):
  93                 if self.response:
  94                         return self.response.headers
  95
  96         def get_header(self, name):
  97                 if self.headers:
  98                         return self.headers.get(name)
  99
 100         @property
 101         def body(self):
 102                 """
 103                         Returns a file-like object with the decoded content
 104                         of the response.
 105                 """
 106                 # Return the response by default
 107                 return self.response