airpac/airpac.py

   1 # -*- coding: utf-8 -*-
   2 # airpac - aria2c wrapper for pacman
   3 # Copyright (C) 2009  Darwin M. Bautista
   4 #
   5 # This program is free software: you can redistribute it and/or modify
   6 # it under the terms of the GNU General Public License as published by
   7 # the Free Software Foundation, either version 3 of the License, or
   8 # (at your option) any later version.
   9 #
  10 # This program is distributed in the hope that it will be useful,
  11 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  12 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13 # GNU General Public License for more details.
  14 #
  15 # You should have received a copy of the GNU General Public License
  16 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
  17
  18 import os
  19 import sys
  20 import shutil
  21 import signal
  22 import asyncore
  23 from tempfile import NamedTemporaryFile
  24 from subprocess import Popen, PIPE
  25
  26
  27 BIN = '/usr/bin/aria2c'
  28 CONF = '/etc/airpac.conf'
  29 STATS = '/var/lib/airpac.stats'
  30 DIR = '/var/lib/pacman/.airpac'
  31 PACDIR = '/var/lib/pacman'
  32
  33 # Obtained from /etc/rc.d/functions
  34 C_MAIN = '\033[1;37;40m'      # main text
  35 C_OTHER = '\033[1;34;40m'     # prefix & brackets
  36 C_BUSY = '\033[0;36;40m'      # busy
  37 C_FAIL = '\033[1;31;40m'      # failed
  38 C_DONE = '\033[1;37;40m'      # completed
  39 C_CLEAR = '\033[1;0m'
  40
  41
  42 def log(width, label, msg, cr=True):
  43     spaces = ' ' * (width - len(label) - len(msg) - 4)
  44     if msg == 'DONE':
  45         status = C_DONE
  46     elif msg == 'FAIL':
  47         status = C_FAIL
  48     else:
  49         status = C_BUSY
  50     print ' %s%s%s%s%s[%s%s%s]%s' % (C_MAIN, label, C_CLEAR, spaces, C_OTHER, status, msg, C_OTHER, C_CLEAR),
  51     if cr:
  52         print '\r',
  53     sys.stdout.flush()
  54
  55
  56 def gen_input_file(uri):
  57     uris = [uri]
  58     max_mirrors = 2
  59     if not uri.endswith('.db.tar.gz'):
  60         with open(CONF, 'r') as conf:
  61             for line in conf:
  62                 if line.lstrip().startswith('split'):
  63                     max_mirrors = 2 * int(line.split('=')[1].strip())
  64         fname = os.path.basename(uri)
  65         base_uri = os.path.dirname(uri)
  66         repo = '[' + uri.split('/')[-4] + ']'
  67         last_repo = None
  68         found_repo = False
  69         with open('/etc/pacman.conf', 'r') as conf:
  70             for line in conf:
  71                 line = line.strip()
  72                 path = line.split('=')[-1].lstrip()
  73                 if line.startswith('['):
  74                     found_repo = (line == repo)
  75                     if not found_repo:
  76                         last_repo = line
  77                     else:
  78                         repo = repo.strip('[]')
  79                         uris = []
  80                 elif found_repo:
  81                     if line.startswith('Server'):
  82                         uris.append('%s/%s\t' % (path, fname))
  83                     elif line.startswith('Include'):
  84                         with open(path, 'r') as f:
  85                             for line in f:
  86                                 line = line.strip()
  87                                 if line.startswith('Server'):
  88                                     path = line.split('=')[-1].lstrip().replace('$repo', repo)
  89                                     uris.append('%s/%s\t' % (path, fname))
  90                                 if len(uris) >= max_mirrors:
  91                                     break
  92                     if len(uris) >= max_mirrors:
  93                         break
  94                 elif not found_repo:
  95                     if base_uri in path and last_repo is not None:
  96                         repo = last_repo
  97                         found_repo = True
  98                         uris.append('\t')
  99     with NamedTemporaryFile(mode='w', prefix='airpac-', delete=False) as tmp:
 100         tmp.writelines(uris)
 101     # Force 2 connections per server if the number of mirrors is less than the value of 'split'
 102     mirrors = len(''.join(uris).split())
 103     split = max_mirrors / 2
 104     if mirrors < split / 2:
 105         split = 2 * mirrors
 106     return tmp.name, str(split)
 107
 108
 109 def db_cache(fname, store=False):
 110     if not os.path.isdir(DIR):
 111         os.mkdir(DIR)
 112     if store:
 113         src = PACDIR
 114         dst = DIR
 115     else:
 116         src = DIR
 117         dst = PACDIR
 118     src += '/' + fname
 119     try:
 120         shutil.copy2(src, dst)
 121     except IOError:
 122         pass
 123
 124
 125 def main():
 126     if not os.access(BIN, os.X_OK):
 127         sys.exit('aria2c not found_repo')
 128     try:
 129         uri, outfile = sys.argv[1:]
 130     except ValueError:
 131         sys.exit('incorrect number of arguments')
 132
 133     outdir = os.path.dirname(outfile)
 134     outfile = os.path.basename(outfile)
 135     tempfile = outfile + '.airpac'
 136     infile, num = gen_input_file(uri)
 137
 138     args = [
 139         BIN, '--conf-path=' + CONF, '--remote-time=true', '--continue',
 140         '--allow-overwrite=true', '--summary-interval=0', '--split=' + num,
 141         '--server-stat-if=' + STATS, '--server-stat-of=' + STATS,
 142         '--dir=' + outdir, '--out=' + tempfile, '--input-file=' + infile
 143     ]
 144
 145     if uri.endswith('.db.tar.gz'):
 146         db_cache(tempfile)
 147
 148     aria2c = Popen(args, stdout=PIPE)
 149
 150     def terminate(signum=None, frame=None, failed=False):
 151         asyncore.socket_map.clear()
 152         aria2c.wait()
 153         os.unlink(infile)
 154         if signum is not None:
 155             print '\n'
 156             sys.stdout.flush()
 157         if signum is not None or failed:
 158             sys.exit(1)
 159
 160     signal.signal(signal.SIGINT, terminate)
 161     signal.signal(signal.SIGTERM, terminate)
 162
 163     # An easy way of getting the terminal width
 164     width = int(Popen(['stty', 'size'], stdout=PIPE).communicate()[0].split()[1])
 165     name = os.path.basename(uri).rsplit('.', 3)[0]
 166
 167     def handle_read():
 168         data = aria2c.stdout.readline().strip()
 169         if data.startswith('[#1'):
 170             log(width, name, data.strip('[#1 ]'))
 171         elif data.startswith('(OK):'):
 172             log(width, name, 'DONE', cr=False)
 173             terminate()
 174         elif data.startswith('(ERR):'):
 175             log(width, name, 'FAIL', cr=False)
 176             terminate(failed=True)
 177
 178     asyncore.file_dispatcher(aria2c.stdout).handle_read = handle_read
 179     asyncore.loop()
 180
 181     if aria2c.returncode == 0:
 182         if uri.endswith('.db.tar.gz'):
 183             db_cache(tempfile, store=True)
 184         os.rename(tempfile, outfile)