#! /usr/bin/python
#
#   rankmirrors - read a list of mirrors from a file and rank them by speed
#   @configure_input@
#
#   Copyright (c) 2002-2007 by Judd Vinet <jvinet@zeroflux.org>
#
#   This program is free software; you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation; either version 2 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program; if not, write to the Free Software
#   Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307,
#   USA.
#
import os, sys, datetime, time, socket, urllib2
from optparse import OptionParser

def createOptParser():
    usage = "usage: %prog [options] MIRRORFILE | URL"
    version = "%prog (pacman) @PACKAGE_VERSION@\n" \
            "Copyright (C) 2002-2007 Judd Vinet <jvinet@zeroflux.org>.\n\n" \
            "This is free software; see the source for copying conditions.\n" \
            "There is NO WARRANTY, to the extent permitted by law."
    description = "Ranks pacman mirrors by their connection and opening " \
            "speed. Pacman mirror files are located in /etc/pacman.d/. It " \
            "can also rank one mirror if the URL is provided."
    parser = OptionParser(usage = usage, version = version,
                          description = description)
    parser.add_option("-n", type = "int", dest = "num", default = 0,
                      help = "number of servers to output, 0 for all")
    parser.add_option("-t", "--times", action = "store_true",
                      dest = "times", default = False,
                      help = "only output mirrors and their response times")
    parser.add_option("-u", "--url", action = "store_true", dest = "url",
                      default = False, help = "test a specific url")
    parser.add_option("-v", "--verbose", action = "store_true",
                      dest = "verbose", default = False,
                      help = "be verbose in ouptut")
    # The following two options should be automatic
    #parser.add_option("-h", "--help", action = "help")
    #parser.add_option("-V", "--version", action = "version")
    return parser

def timeCmd(cmd):
    before = time.time()
    try:
        cmd()
    except KeyboardInterrupt, ki:
        raise ki
    except socket.timeout, ioe:
        return 'timeout'
    except Exception, e:
        return 'unreachable'
    return time.time() - before

def talkToServer(serverUrl):
    opener = urllib2.build_opener()
    # retrieve first 50,000 bytes only
    tmp = opener.open(serverUrl).read(50000)

def getFuncToTime(serverUrl):
    return lambda : talkToServer(serverUrl)

def cmpPairBySecond(p1, p2):
    if p1[1] == p2[1]:
        return 0
    if p1[1] < p2[1]:
        return -1
    return 1

def printResults(servers, time, verbose, num):
    items = servers.items()
    items.sort(cmpPairBySecond)
    itemsLen = len(items)
    numToShow = num
    if numToShow > itemsLen or numToShow == 0:
        numToShow = itemsLen
    if itemsLen > 0:
        if time:
            print
            print ' Servers sorted by time (seconds):'
            for i in items[0:numToShow]:
                if i[1] == 'timeout' or i[1] == 'unreachable':
                    print i[0], ':', i[1]
                else:
                    print i[0], ':', "%.2f" % i[1]
        else:
            for i in items[0:numToShow]:
                print 'Server =', i[0]

if __name__ == "__main__":
    parser = createOptParser()
    (options, args) = parser.parse_args()

    if len(args) != 1:
        parser.print_help(sys.stderr)
        sys.exit(0)

    # allows connections to time out if they take too long
    socket.setdefaulttimeout(10)

    if options.url:
        if options.verbose:
            print 'Testing', args[0] + '...'
        try:
            serverToTime = timeCmd(getFuncToTime(args[0]))
        except KeyboardInterrupt, ki:
            sys.exit(1)
        if serverToTime == 'timeout' or serverToTime == 'unreachable':
            print args[0], ':', serverToTime
        else:
            print args[0], ':', "%.2f" % serverToTime
        sys.exit(0)

    if not os.path.isfile(args[0]):
        print >>sys.stderr, 'rankmirrors: file', args[0], 'does not exist.'
        sys.exit(1)

    fl = open(args[0], 'r')
    serverToTime = {}
    if options.times:
        print 'Querying servers, this may take some time...'
    else:
        print "# Server list generated by rankmirrors on",
        print datetime.date.today()
    for ln in fl.readlines():
        splitted = ln.split('=')
        if splitted[0].strip() != 'Server':
            if not options.times:
                print ln,
            continue

        serverUrl = splitted[1].strip()
        if serverUrl[-1] == '\n':
            serverUrl = serverUrl[0:-1]
        if options.verbose and options.times:
            print serverUrl, '...',
        elif options.verbose:
            print '#', serverUrl, '...',
        elif options.times:
            print ' * ',
        sys.stdout.flush()

        # add *.db.tar.gz to server name. the repo name is parsed
        # from the mirror url; it is the third (or fourth) dir
        # from the end, where the url is http://foo/bar/REPO/os/arch
        try:
            splitted2 = serverUrl.split('/')
            if serverUrl[-1] != '/':
                repoName = splitted2[-3]
                dbFileName = '/' + repoName + '.db.tar.gz'
            else:
                repoName = splitted2[-4]
                dbFileName = repoName + '.db.tar.gz'
        except:
            dbFileName = ''

        try:
            serverToTime[serverUrl] = timeCmd(getFuncToTime(serverUrl + dbFileName))
            if options.verbose:
                try:
                    print "%.2f" % serverToTime[serverUrl]
                except:
                    print serverToTime[serverUrl]
        except:
            print
            printResults(serverToTime, options.times, options.verbose,
                         options.num)
            sys.exit(0)

    printResults(serverToTime, options.times, options.verbose, options.num)

# vim: set ts=4 sw=4 et: