#$Id: upgrade.py 2192 2005-05-12 09:57:20Z joag $

import threading
import sys
import os
import logging
import subprocess
import string
import time
import ConfigParser
import atexit
import tempfile
from optparse import OptionParser


def cleanup():
    """nicely cleanup everthing before exiting"""
    logging.shutdown()

# register a function call to cleanup stuff before exiting
atexit.register(cleanup)


class Tftp(threading.Thread):
    """download a file with tftp to an address"""
    def __init__(self, address, firmware, destination, coldboot, coldboottimeout):
        threading.Thread.__init__(self)
        self.address = address
        self.firmware = firmware
        self.destination = destination
        self.coldboot = coldboot
        self.coldboottimeout = coldboottimeout
    def run(self):
        filelogger.info("starting tftp to '%s' of '%s' --> '%s'" \
                        %(self.address, self.firmware, self.destination))
        p = subprocess.Popen('tftp -i ' + self.address + ' PUT "' + \
                             self.firmware + '" "' + self.destination + '"', \
                             stdin=None, stdout=subprocess.PIPE, \
                             stderr=subprocess.STDOUT, universal_newlines=True)
        # this can take a while
        returncode = p.wait()
        errormessage = p.stdout.readlines()
        filelogger.info("tftp to '%s' finished with return code '%i'" \
                        %(self.address, returncode))
        for line in errormessage:
            filelogger.debug("tftp to '%s' output '%s'" \
                             %(self.address, string.rstrip(line, '\n')))
        if returncode != 0:
            for line in errormessage:
                line = string.rstrip(line, '\n')
                if (len(line) != 0):
                    filelogger.error("tftp to '%s' errormessage '%s'" \
                             %(self.address, line))
        else:
            if self.coldboot == True:
                # sleep for the specified timeout (for SHDSL)
                filelogger.info("sleeping %i seconds for '%s' before coldboot" 
                                    %(self.coldboottimeout, self.address))
                time.sleep(self.coldboottimeout)
                filelogger.info("doing coldboot on '%s'" %self.address)
                output = tmaclicommand(self.address, '-action "Cold Boot"')
                for line in output:
                    line = string.rstrip(line, '\n')
                    if len(line) != 0:
                        if string.find(line,'NOK') != -1:
                            filelogger.error("Coldboot to '%s' failed" %self.address)
                            break
        # when this function returns, the thread finishes


class Tml(threading.Thread):
    """download a file with tml to an address"""
    def __init__(self, address, firmware, destination, coldboot, coldboottimeout):
        threading.Thread.__init__(self)
        self.address = address
        self.firmware = firmware
        self.destination = destination
        self.coldboot = coldboot
        self.coldboottimout = coldboottimeout
    def run(self):
        filelogger.info("starting tml to '%s' of '%s' --> '%s'" \
                        %(self.address, self.firmware, self.destination))
        # create tempfilename for logging
        logfilename = tempfile.mkstemp()
        os.close(logfilename[0])
        p = subprocess.Popen('tml -z' + logfilename[1] + ' ' + self.address + ' -f"' + \
                            self.firmware + '"@"' + self.destination + '"', \
                            stdin=None, stdout=subprocess.PIPE, \
                             stderr=subprocess.STDOUT, universal_newlines=True)
        # this can take a while
        returncode = p.wait()
        # open the tempfile containing the logs of tml
        logfile = open(logfilename[1])
        filelogger.info("tml to '%s' finished with return code '%i'" \
                        %(self.address, returncode))
        for line in logfile:
            filelogger.debug("tml to '%s' log '%s'" \
                             %(self.address, string.rstrip(line, '\n')))
        if returncode != 0:
            for line in logfile:
                if (len(line) != 0):
                    filelogger.error("tml to '%s' errormessage '%s'" \
                             %(self.address, string.rstrip(line, '\n')))
        else:
            if self.coldboot == True:
                # sleep for the specified timeout (for SHDSL)
                time.sleep(self.coldboottimeout)
                filelogger.debug("doing coldboot on '%s'" %self.address)
                output = tmaclicommand(self.address, '-action "Cold Boot"')
                for line in output:
                    line = string.rstrip(line, '\n')
                    if len(line) != 0:
                        if string.find(line,'NOK') != -1:
                            filelogger.error("Coldboot to '%s' failed" %self.address)
                            break
        # clean up the temp logfile
        logfile.close()
        os.remove(logfilename[1])
        # when this function returns, the thread finishes


def tmaclicommand(address, command):
    tmaclitje = subprocess.Popen('tmacli ' + address + ' ' + command,
                         stdin=None, stdout=subprocess.PIPE,
                         stderr=subprocess.STDOUT, universal_newlines=True)
    returncode = tmaclitje.wait()
    output = tmaclitje.stdout.readlines()
    return output

def count_active(threads):
    """returns the number threads that are alive"""
    num_active = 0
    for thread in threads:
        if thread.isAlive():
            num_active += 1
    #print '%d alive' % num_active
    return num_active

def main():
    # options parsing
    usage = "usage: %prog [options] filename"
    parser = OptionParser(usage, 
            version="%prog $Id: upgrade.py 2192 2005-05-12 09:57:20Z joag $")
    parser.add_option("-v", "--verbose", action="count", dest="verbosity", 
                      help="be more verbose, "
                           "verbosity can be increased to 3 "
                           "[default=%default]")
    parser.add_option("-q", "--quiet", action="store_true", dest="quiet",
                      help="don't even print errormessages")
    parser.set_defaults(quiet=False)
    parser.set_defaults(verbosity=1)

    (options, args) = parser.parse_args()

    if options.quiet == True:
        options.verbosity = 0
    if (options.verbosity > 2):
        print "opions: %s" %options
        print "args %s" %args

    if len(args) != 1:
        if (options.verbosity > 0):
            parser.error("please specify one filename")
    filename = args[0]


    INIFILE = 'upgrade.ini'
    # read configs from INIFILE
    if os.path.isfile(INIFILE):
        config = ConfigParser.ConfigParser()
        if (options.verbosity > 1):
            print "reading '%s'" % INIFILE
        config.read(INIFILE)
        MAX_THREADS = int(config.get('upgrade', 'MAX_THREADS'))
        LOGFILE = config.get('upgrade', 'LOGFILE')
        # The default levels provided are DEBUG, INFO, WARNING, ERROR and CRITICAL.
        LOGLEVEL = config.get('upgrade', 'LOGLEVEL')
    else:
        sys.exit('could not open file \'%s\' in dir \'%s\'' %(INIFILE, os.getcwd()))

    # configure the logger
    global filelogger
    filelogger = logging.getLogger('upgrade')
    hdlr = logging.FileHandler(LOGFILE)
    hdlr2 = logging.StreamHandler(sys.stdout)
    formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s')
    hdlr.setFormatter(formatter)
    hdlr2.setFormatter(formatter)
    filelogger.addHandler(hdlr)
    filelogger.setLevel(int(eval('logging.' + LOGLEVEL)))
    
    # log some variables for debugging
    filelogger.debug('MAX_THREADS = %d' %MAX_THREADS)
    filelogger.debug('LOGFILE = %s' %LOGFILE)
    filelogger.debug('logging with level %s' % logging.getLevelName(filelogger.getEffectiveLevel()))
    
    # open the file with the tasks
    try:
        f = open(filename)
        if (options.verbosity > 1):
            print "reading '%s'" % filename
        buffer = f.read()
        lines = buffer.splitlines()
        f.close()
    except:
        filelogger.critical("could not open file '%s'" % sys.argv[1])
        sys.exit("could not open file '%s'" %sys.argv[1])
    
    # read the lines, and filter out the tasks
    tasks = []
    for i in range(0,len(lines)):
        # whitespace (strip) and lines starting with # are comment
        if (len(string.strip(lines[i])) != 0 and \
                string.strip(lines[i])[0] != '#'):
            # tab separated values
            arguments = lines[i].split('\t')
            if len(arguments) >= 4:
                method = arguments[0]
                address = arguments[1]
                firmware = arguments[2]
                destination = arguments[3]
                if len(arguments) >= 6:
                    coldboot = bool(arguments[4]=='True')
                    coldboottimeout = int(arguments[5])
                else:
                    coldboot = False
                    coldboottimeout = None

                filelogger.info("method='%s' address='%s' firmware='%s' destination='%s' coldboot='%s' coldboottimeout='%s'" \
                                %(method, address, firmware, destination, coldboot, coldboottimeout))
                tasks.append({'method':method, 'address':address, \
                              'firmware':firmware, 'destination':destination, \
                            'coldboot':coldboot, 'coldboottimeout':coldboottimeout})
            else:
                filelogger.error("Linenr '%d' needs at least 4 arguments" %(i))
                break

        else:
            filelogger.debug("Linenr '%d', this is comment '%s'" %(i, lines[i]))
    filelogger.info("'%d' tasks found" %(len(tasks)))
    
    # start the downloaders, respecting MAX_THREADS
    downloaders = []
    for task in tasks:
        while count_active(downloaders) >= MAX_THREADS:
            #print ' too many active'
            time.sleep(1)
        # instanciate the respective class with the correct arguments
        downloader = eval(task['method'])(task['address'], task['firmware'], \
                                       task['destination'], task['coldboot'], \
                                       task['coldboottimeout'])
        downloaders.append(downloader)
        downloader.start()
    
    # wait for all the downloaders to finish
    for downloader in downloaders:
        # join() causes the main thread to wait until downloader is finished
        downloader.join()
    
    filelogger.info("'%d' downloads finished" %(len(downloaders)))

if __name__ == "__main__":
    main()