Cnchi/rank_mirrors.py

   1 #!/usr/bin/env python
   2 # -*- coding: utf-8 -*-
   3 #
   4 # rank_mirrors.py
   5 #
   6 # Copyright © 2012, 2013 Xyne
   7 # Copyright © 2013-2018 Antergos
   8 #
   9 # This file is part of Cnchi.
  10 #
  11 # Cnchi is free software; you can redistribute it and/or modify
  12 # it under the terms of the GNU General Public License as published by
  13 # the Free Software Foundation; either version 3 of the License, or
  14 # (at your option) any later version.
  15 #
  16 # Cnchi is distributed in the hope that it will be useful,
  17 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  18 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  19 # GNU General Public License for more details.
  20 #
  21 # The following additional terms are in effect as per Section 7 of the license:
  22 #
  23 # The preservation of all legal notices and author attributions in
  24 # the material or in the Appropriate Legal Notices displayed
  25 # by works containing it is required.
  26 #
  27 # You should have received a copy of the GNU General Public License
  28 # along with Cnchi; If not, see <http://www.gnu.org/licenses/>.
  29
  30
  31 """ Creates mirrorlist sorted by both latest updates and fastest connection """
  32
  33 import http.client
  34 import logging
  35 import multiprocessing
  36 import os
  37 import queue
  38 import subprocess
  39 import threading
  40 import time
  41 import urllib.request
  42 import urllib.error
  43
  44 import feedparser
  45 import requests
  46
  47 import update_db
  48 import misc.extra as misc
  49
  50 # When testing, no _() is available
  51 try:
  52     _("")
  53 except NameError as err:
  54     def _(message):
  55         return message
  56
  57 class RankMirrors(multiprocessing.Process):
  58     """ Process class that downloads and sorts the mirrorlist """
  59
  60     REPOSITORIES = ['arch']
  61     MIRROR_OK_RSS = 'Alert Details: Successful response received'
  62
  63     MIRROR_STATUS = {
  64         'arch': 'http://www.archlinux.org/mirrors/status/json/'}
  65
  66     MIRRORLIST = {
  67         'arch': '/etc/pacman.d/mirrorlist'}
  68
  69     MIRRORLIST_URL = {
  70         'arch': "https://www.archlinux.org/mirrorlist/all/"}
  71
  72     DB_SUBPATHS = {
  73         'arch': 'core/os/x86_64/{0}-{1}-x86_64.pkg.tar.xz'}
  74
  75     def __init__(self, fraction_pipe, settings):
  76         """ Initialize process class
  77             fraction_pipe is a pipe used to send progress for a gtk.progress widget update
  78             in another process (see start_rank_mirrors() in mirrors.py) """
  79         super(RankMirrors, self).__init__()
  80         self.settings = settings
  81         self.fraction_pipe = fraction_pipe
  82         # Antergos mirrors info is returned as RSS, arch's as JSON
  83         self.data = {'arch': {}}
  84         self.mirrorlist_ranked = {'arch': []}
  85
  86     @staticmethod
  87     def is_good_mirror(mirror):
  88         """ Check if mirror info is good enough """
  89         if 'summary' in mirror.keys():
  90             # RSS antergos status mirror
  91             return bool(mirror['summary'] == RankMirrors.MIRROR_OK_RSS)
  92
  93         # JSON arch status mirror
  94         return (mirror['last_sync'] and
  95                 mirror['completion_pct'] == 1.0 and
  96                 mirror['protocol'] == 'http' and
  97                 int(mirror['delay']) <= 3600)
  98
  99     def get_mirror_stats(self):
 100         """ Retrieve all mirrors status RSS data. """
 101         # Load status data (JSON) for arch mirrors
 102         if not self.data['arch']:
 103             try:
 104                 req = requests.get(
 105                     RankMirrors.MIRROR_STATUS['arch'],
 106                     headers={'User-Agent': 'Mozilla/5.0'}
 107                 )
 108                 self.data['arch'] = req.json()
 109             except requests.RequestException as err:
 110                 logging.warning(
 111                     'Failed to retrieve mirror status information: %s', err)
 112
 113         # Load status data (RSS) for antergos mirrors
 114 #        if not self.data['antergos']:
 115 #            self.data['antergos'] = feedparser.parse(
 116 #                RankMirrors.MIRROR_STATUS['antergos'])
 117
 118         mirrors = {'arch': []}
 119
 120         try:
 121             # Filter incomplete mirrors and mirrors that haven't synced.
 122             mirrors['arch'] = self.data['arch']['urls']
 123             mirrors['arch'] = [m for m in mirrors['arch'] if self.is_good_mirror(m)]
 124             #self.data['arch']['urls'] = mirrors['arch']
 125         except KeyError as err:
 126             logging.warning('Failed to parse retrieved mirror data: %s', err)
 127
 128         mirror_urls = []
 129 #        for mirror in self.data['antergos']['entries']:
 130 #            title = mirror['title']
 131 #            if "is UP" in title:
 132 #                # In RSS, all mirrors are in http:// format, we prefer https://
 133 #                mirror['url'] = mirror['link'].replace('http://', 'https://')
 134 #                if mirror['url'] not in mirror_urls:
 135 #                    mirrors['antergos'].append(mirror)
 136 #                    mirror_urls.append(mirror['url'])
 137
 138         return mirrors
 139
 140     @staticmethod
 141     def get_antergos_mirror_url(mirror_url):
 142         """ Get full mirror url from the stats mirror url """
 143         lines = []
 144         mirrorlist_path = RankMirrors.MIRRORLIST['antergos']
 145         with open(mirrorlist_path, 'r') as mirror_file:
 146             lines = mirror_file.readlines()
 147         for line in lines:
 148             if mirror_url in line:
 149                 return line.split('=')[1].strip()
 150         logging.warning("%s not found in %s", mirror_url, mirrorlist_path)
 151         return None
 152
 153     @staticmethod
 154     def get_package_version(name):
 155         """ Returns pkg_name package version """
 156         try:
 157             cmd = ["/usr/bin/pacman", "-Ss", name]
 158             line = subprocess.check_output(cmd).decode().split()
 159             version = line[1]
 160             logging.debug(
 161                 '%s version is: %s (used to test mirror speed)', name, version)
 162         except subprocess.CalledProcessError as err:
 163             logging.warning(err)
 164             version = False
 165         return version
 166
 167     def sort_mirrors_by_speed(self, mirrors=None, max_threads=8):
 168         """ Sorts mirror list """
 169
 170         test_packages = {
 171             'arch': {'name':'cryptsetup', 'version': ''}}
 172
 173         rated_mirrors = {'arch': []}
 174
 175         for key, value in test_packages.items():
 176             test_packages[key]['version'] = self.get_package_version(value['name'])
 177
 178         total_num_mirrors = 0
 179         for key in mirrors.keys():
 180             total_num_mirrors += len(mirrors[key])
 181         num_mirrors_done = 0
 182         old_fraction = -1
 183
 184         num_threads = min(max_threads, total_num_mirrors)
 185         # URL input queue.Queue
 186         q_in = queue.Queue()
 187         # URL and rate output queue.Queue
 188         q_out = queue.Queue()
 189
 190         name = ""
 191         version = ""
 192         rates = {}
 193
 194         for repo in RankMirrors.REPOSITORIES:
 195             name = test_packages[repo]['name']
 196             version = test_packages[repo]['version']
 197
 198             def worker():
 199                 """ worker thread. Retrieves data to test mirror speed """
 200                 while not q_in.empty():
 201                     mirror_url, full_url = q_in.get()
 202                     # Leave the rate as 0 if the connection fails.
 203                     rate = 0
 204                     dtime = float('NaN')
 205                     if full_url:
 206                         req = urllib.request.Request(url=full_url)
 207                         try:
 208                             time0 = time.time()
 209                             with urllib.request.urlopen(req, None, 5) as my_file:
 210                                 size = len(my_file.read())
 211                                 dtime = time.time() - time0
 212                                 rate = size / dtime
 213                         except (OSError, urllib.error.HTTPError,
 214                                 http.client.HTTPException) as err:
 215                             logging.warning("Couldn't download %s", full_url)
 216                             logging.warning(err)
 217                     q_out.put((mirror_url, full_url, rate, dtime))
 218                     q_in.task_done()
 219
 220             # Load the input queue.Queue
 221             url_len = 0
 222             for mirror in mirrors[repo]:
 223                 url_len = max(url_len, len(mirror['url']))
 224                 if repo == 'antergos':
 225                     url = self.get_antergos_mirror_url(mirror['url'])
 226                     # Save mirror url
 227                     mirror['url'] = url
 228                     if url is None:
 229                         package_url = None
 230                     else:
 231                         # Compose package url
 232                         package_url = url.replace('$repo', 'antergos').replace('$arch', 'x86_64')
 233                         package_url += RankMirrors.DB_SUBPATHS['antergos'].format(name, version)
 234                 else:
 235                     package_url = mirror['url']
 236                 if mirror['url'] and package_url:
 237                     q_in.put((mirror['url'], package_url))
 238
 239             # Launch threads
 240             my_threads = []
 241             for _index in range(num_threads):
 242                 my_thread = threading.Thread(target=worker)
 243                 my_thread.start()
 244                 my_threads.append(my_thread)
 245
 246             # Remove mirrors that are not present in antergos-mirrorlist
 247             if repo == 'antergos':
 248                 mirrors_pruned = []
 249                 for mirror in mirrors[repo]:
 250                     if mirror['url'] is not None:
 251                         mirrors_pruned.append(mirror)
 252                 mirrors[repo] = mirrors_pruned
 253
 254             # Wait for queue to empty
 255             while not q_in.empty():
 256                 fraction = (float(q_out.qsize()) + num_mirrors_done) / float(total_num_mirrors)
 257                 if fraction != old_fraction:
 258                     if self.fraction_pipe:
 259                         self.fraction_pipe.send(fraction)
 260                     old_fraction = fraction
 261
 262             num_mirrors_done += q_out.qsize()
 263
 264             # Wait for all threads to complete
 265             q_in.join()
 266
 267             # Log some extra data.
 268             url_len = str(url_len)
 269             fmt = '%-' + url_len + 's  %14s  %9s'
 270             logging.debug(fmt, _("Server"), _("Rate"), _("Time"))
 271
 272             # Loop over the mirrors just to ensure that we get the rate for each.
 273             # The value in the loop does not (necessarily) correspond to the mirror.
 274             fmt = '%-' + url_len + 's  %8.2f KiB/s  %7.2f s'
 275             for mirror in mirrors[repo]:
 276                 url, full_url, rate, dtime = q_out.get()
 277                 if full_url:
 278                     kibps = rate / 1024.0
 279                     logging.debug(fmt, url, kibps, dtime)
 280                     rates[url] = rate
 281                 q_out.task_done()
 282
 283             # Wait for all threads to finnish (all will be finished, but...)
 284             for my_thread in my_threads:
 285                 my_thread.join()
 286
 287             # Sort mirrors by rate
 288             try:
 289                 rated_mirrors[repo] = [m for m in mirrors[repo] if rates[m['url']] > 0]
 290                 rated_mirrors[repo].sort(key=lambda m: rates[m['url']], reverse=True)
 291             except KeyError as err:
 292                 logging.warning(err)
 293
 294         return rated_mirrors
 295
 296     @staticmethod
 297     def uncomment_mirrors():
 298         """ Uncomment mirrors and comment out auto selection so
 299         rankmirrors can find the best mirror. """
 300
 301 #        comment_urls = [
 302 #            'http://mirrors.antergos.com/$repo/$arch',
 303 #            'sourceforge']
 304
 305         for repo in RankMirrors.REPOSITORIES:
 306             if os.path.exists(RankMirrors.MIRRORLIST[repo]):
 307                 with open(RankMirrors.MIRRORLIST[repo]) as mirrors:
 308                     lines = [x.strip() for x in mirrors.readlines()]
 309
 310                 for i, line in enumerate(lines):
 311                     if line.startswith("#Server"):
 312                         # if server is commented, uncoment it.
 313                         lines[i] = line.lstrip("#")
 314
 315                     if line.startswith("Server"):
 316                         # Let's see if we have to comment out this server
 317                         for url in comment_urls:
 318                             if url in line:
 319                                 lines[i] = '#' + line
 320
 321                 # Write new one
 322                 with misc.raised_privileges():
 323                     try:
 324                         with open(RankMirrors.MIRRORLIST[repo], 'w') as mirrors_file:
 325                             mirrors_file.write("\n".join(lines) + "\n")
 326                     except (OSError, PermissionError) as err:
 327                         logging.error(err)
 328         update_db.sync()
 329
 330     def filter_and_sort_mirrorlists(self):
 331         """ Filter and sort mirrors """
 332
 333         mlist = self.get_mirror_stats()
 334         mirrors = self.sort_mirrors_by_speed(mirrors=mlist)
 335
 336         for repo in ['arch']:
 337             self.mirrorlist_ranked[repo] = []
 338
 339         for repo in ['arch']:
 340             output = '# {} mirrorlist generated by cnchi #\n'.format(repo)
 341             for mirror in mirrors[repo]:
 342                 self.mirrorlist_ranked[repo].append(mirror['url'])
 343                 if repo == 'arch':
 344                     output += "Server = {0}{1}/os/{2}\n".format(mirror['url'], '$repo', '$arch')
 345                 else:
 346                     output += "Server = {0}\n".format(mirror['url'])
 347
 348             # Write modified mirrorlist
 349             with misc.raised_privileges():
 350                 try:
 351                     with open(RankMirrors.MIRRORLIST[repo], 'w') as mirrors_file:
 352                         mirrors_file.write(output)
 353                 except (OSError, PermissionError) as err:
 354                     logging.error(err)
 355                 update_db.sync()
 356
 357     def run(self):
 358         """ Run process """
 359         # Wait until there is an Internet connection available
 360         while not misc.has_connection():
 361             time.sleep(2)  # Delay, try again after 2 seconds
 362
 363         logging.debug("Updating both mirrorlists (Arch and Antergos)...")
 364         self.update_mirrorlists()
 365
 366         self.uncomment_mirrors()
 367
 368         logging.debug("Filtering and sorting mirrors...")
 369         self.filter_and_sort_mirrorlists()
 370
 371         if self.settings:
 372             self.mirrorlist_ranked['arch'] = [
 373                 x for x in self.mirrorlist_ranked['arch'] if x]
 374             self.settings.set('rankmirrors_result', self.mirrorlist_ranked['arch'])
 375
 376         if self.fraction_pipe:
 377             self.fraction_pipe.send(1)
 378             self.fraction_pipe.close()
 379
 380         logging.debug("Auto mirror selection has been run successfully.")
 381
 382
 383     @staticmethod
 384     def update_mirrorlists():
 385         """ Download mirror lists from archlinux and github """
 386         for repo in RankMirrors.REPOSITORIES:
 387             url = RankMirrors.MIRRORLIST_URL[repo]
 388             req = urllib.request.Request(url=url)
 389             try:
 390                 with urllib.request.urlopen(req, None, 5) as my_file:
 391                     data = my_file.read()
 392                 with misc.raised_privileges():
 393                     with open(RankMirrors.MIRRORLIST[repo], 'wb') as mirror_file:
 394                         mirror_file.write(data)
 395             except (OSError, urllib.error.HTTPError, http.client.HTTPException) as err:
 396                 logging.warning("Couldn't download %s", url)
 397                 logging.warning(err)
 398
 399 def test_module():
 400     """ Helper function to test this module """
 401     logger = logging.getLogger()
 402     logger.setLevel(logging.DEBUG)
 403
 404     proc = RankMirrors(None, None)
 405     proc.daemon = True
 406     proc.name = "rankmirrors"
 407     proc.start()
 408     proc.join()
 409
 410 if __name__ == '__main__':
 411     test_module()