#!/usr/bin/env python

import os,sys,commands
import pycurl
import cStringIO
import debconf
import operator
import re

MIRROR_LIST_FILE="/usr/share/linuxmint/mint-debian-mirrors/mirrors.list"

class PostInst:
    def main(self):                
        # this script can get called multiple times
        # so we only set up the frontend once
        # thanks to the mailman developers for this section
        try:
            os.environ['DEBIAN_HAS_FRONTEND']
        except:
            # we need to invoke the frontend ourselves as
            # the debconf module does not strip the python cruft off
            os.execv(debconf._frontEndProgram, [debconf._frontEndProgram] + sys.argv)

        db = debconf.Debconf()

        if len(sys.argv) > 1:
            if sys.argv[1] == 'configure':
                dry_run = False
                if (len(sys.argv) > 2 and sys.argv[2] == "--dry-run"):
                    dry_run = True
                try:
                    if not os.path.exists(MIRROR_LIST_FILE):
                        print "Can't find mirror list file"
                        sys.exit(0)

                    if os.getuid() != 0:
                        print "Warning : script launched in user mode : can't modify anything"
                        sys.exit(0)
                        
                    current_server = None
                    current_repo = "latest"
                    deb_lines = [line.strip() for line in open("/etc/apt/sources.list")]                                    
                    for deb_line in deb_lines:      
                        deb_line = str(unicode(deb_line, 'ascii', 'ignore'))
                        if deb_line.lstrip().startswith("#"):                            
                            continue
                        elements = deb_line.split(" ")
                        if ("debian.org" in deb_line or "debian.net" in deb_line):
                            print "Found debian.org in sources.list. Skipping mirror selection."
                            sys.exit(0)                            
                        if ("testing/updates main contrib non-free" in deb_line):                                                        
                            if ("http://" in elements[1] or "ftp://" in elements[1]):
                                current_server = elements[1]
                            elif ("http://" in elements[2] or "ftp://" in elements[2]):
                                current_server = elements[2]
                            elif ("http://" in elements[3] or "ftp://" in elements[3]):
                                current_server = elements[3]
                            else:
                                print "Can't find current server info. Skipping."
                                sys.exit(0)
                            current_server = current_server.replace("/security", "")
                            current_server = current_server.replace("/latest", "")
                            if ("incoming" in current_server):
                                current_repo = "incoming"
                                current_server = current_server.replace("/incoming", "")
                                        
                    if current_server is None:
                        print "Current server not found, exiting"
                        sys.exit(0)                        
                    else:
                        current_server = current_server.strip()   
                        current_server = current_server.replace('\t', '')                   
                        current_server = current_server.replace('\n', '')                                                                     
                        if current_server.endswith('/'):
                            current_server = current_server[:-1]
                        print "Current server: '%s' (%s)" % (current_server, current_repo)

                    mirrors = [line.strip() for line in open(MIRROR_LIST_FILE)]

                    best_speed = 0
                    best_server = None
                    choices = []

                    for mirror in mirrors:
                        mirror = str(unicode(mirror, 'ascii', 'ignore'))
                        if ("http" in mirror or "ftp" in mirror):        
                            if mirror.endswith('/'):
                                mirror = mirror[:-1]
                            c = pycurl.Curl()
                            buff = cStringIO.StringIO()
                            c.setopt(pycurl.URL, "%s/latest/dists/testing/Release" % mirror)
                            c.setopt(pycurl.CONNECTTIMEOUT, 10)
                            c.setopt(pycurl.TIMEOUT, 10)
                            c.setopt(pycurl.FOLLOWLOCATION, 1)        
                            c.setopt(pycurl.WRITEFUNCTION, buff.write)
                            try:
                              c.perform()
                            except pycurl.error,error:
                              errno, errstr = error
                              print 'An error occurred: ', errstr
                              continue
                            return_code = c.getinfo(pycurl.HTTP_CODE)
                            download_speed = c.getinfo(pycurl.SPEED_DOWNLOAD)
                            if (return_code == 200):                         
                                download_speed = int(round(download_speed/1000))
                                choices.append([mirror,download_speed])
                                print "Server %s - %dKbps" % (mirror, download_speed)
                                if (download_speed > best_speed):
                                    best_speed = download_speed
                                    best_server = mirror
                            else:
                                print "Warning %s/latest/dists/testing/Release returns HTTP code %d" % (mirror, return_code)
                            
                    print "Best server: %s" % best_server   
                    
                    if (dry_run == True):
                        sys.exit(0)

                    if best_server is not None: # and current_server != best_server:    
                        # Sort choices, best first                        
                        choices = sorted(choices, key=operator.itemgetter(1), reverse=True)
                        formatted_choices=""
                        for choice in choices:
                          formatted_choices = '%s%s - %dKbps, ' % (formatted_choices, choice[0], choice[1])
                        formatted_choices = formatted_choices[:-2]
                        db.fset('mint-debian-mirrors/choose-mirror','seen','false')
                        db.subst('mint-debian-mirrors/choose-mirror', "choices", formatted_choices)                        
                        db.subst('mint-debian-mirrors/choose-mirror', "current_server", current_server)                        
                        db.input(debconf.CRITICAL, 'mint-debian-mirrors/choose-mirror')                                  
                        db.go()
                        selected_server = db.get('mint-debian-mirrors/choose-mirror')                        
                        selected_server=selected_server.split(' ')[0]
                        print "Selected server: %s" % selected_server
                        # Substitude in sources.list                        
                        db.unregister('mint-debian-mirrors/choose-mirror')
                        if (selected_server is not None and selected_server != ""):
                            print "Saving old /etc/apt/sources.list in /etc/apt/sources.list.bk"
                            os.system("cp /etc/apt/sources.list /etc/apt/sources.list.bk")
                            print "Updating /etc/apt/sources.list"
                            
                            new_deb_lines = []                                                    
                            for deb_line in deb_lines:  
                                deb_line = str(unicode(deb_line, 'ascii', 'ignore'))
                                new_deb_line = deb_line
                                if deb_line.lstrip().startswith("#"):
                                    new_deb_lines.append(new_deb_line)
                                    continue                                
                                if current_server in deb_line:
                                    elements = deb_line.split(" ")
                                    if ("latest" in deb_line or "incoming" in deb_line):
                                        new_deb_line = deb_line.replace(current_server, selected_server)
                                    else:
                                        new_deb_line = deb_line.replace(current_server, "%s/%s" % selected_server, current_repo)                                    
                                new_deb_lines.append(new_deb_line)
                                                                                                                                                
                            f = open('/etc/apt/sources.list', 'w')
                            for deb_line in new_deb_lines:
                                f.write("%s\n" % deb_line)
                            f.close()

                except Exception, detail:
                    # This is a best-effort attempt, fail graciously
                    print detail                     
                
            elif sys.argv[1] == "reconfigure":
                previous_version = sys.argv[2]
                # we're in reconfigure mode
                
if __name__ == "__main__":
    PostInst().main()

