# xbu_upgr.py (v0.3.1)
# upgrade class for XB firmware
#
# 9-18-2010 NWF
#
# TODO: -add fwupgr list for all types of fw and support
#        to check and upgrade each based on firmware type + ver
#       -add different __process methods for other methods of checking
#       -add DTD capabilities
#       -add ability to upgrade gateway fw
# 
import os
import sys
import string
import time
import threading
import base64
import datetime
import threading
from msvcrt import getch
from Queue import Queue
from time import strftime
#import xml.sax
#from xml.sax import handler, make_parser, parseString
import idigi_comm
import idigi_parse

# connStatus offsets
CS_DISCONNECTED = 0
CS_CONNECTED = 1

# firmware upgrade status constants
US_STOPPED = 0
US_UPGRADING = 1
US_SUCCESS_OK = 2
US_SUCCESS_NOT_OK = 3

# master device list offsets
ML_CWID = 0
ML_MAC = 1
ML_NAME = 2
ML_CONSTAT = 3
ML_ASSOC = 4
ML_HW = 5
ML_FW = 6

# fw_type constants
ZNET_FW = "1"
ZB_FW = "2"
SMARTENERGY_FW = "3"
DIGIMESH_FW = "8"

# task offsets
TSK_TASKTYPE = 0
TSK_CWID = 1
TSK_MAC = 2
TSK_ORIG = 3
TSK_CURR = 4
TSK_STAT = 5
TSK_CALLBACK = 6
TSK_FWTYPE = 7
TSK_TIME = 8

radio_types = { "0x19":"S2", "0x1a":"S2Pro", "0x1e": "S2B Pro" }

def DumpToFile(xmlfilename, xmldata):
    filenamepath = "./" + xmlfilename
    print "DumpToFile: filepath= ", xmlfilename
    if (os.path.exists(filenamepath)):    
        handle = open(filenamepath, 'a')
    else:
        handle = open(filenamepath, 'w')
    handle.write(xmldata)
    handle.close()
    return True

class IDigiDeviceDict(dict):
    def __init__(self, default='na'):
        print "\n\nIDigiDeviceDict:"
        print "\n\nIDigiDeviceDict template"
        self.default = default
        dict.__init__(self)
        
    def __getitem__(self, key):
        try:
            print "\n\nIDigiDeviceDict __getitem__"
            return dict.__getitem__(self, key)

        except KeyError:
            print "IDIGIDEV: KeyError"
            return self.default
        
class Logger:
    def __init__(self, logfilepath):
        self.logfilepath = logfilepath
        self.filehandle = None

## task = type=0, cwid=1, mac=2, orig-3, curfw=4, status=5, cb=6,
##        fwtype=7, timestamp=8 
## task = ["UPGR", key, origfw, curfw, reason, TaskComplRoutine, \
##         ZB_FW,] # timestamp to be appended by complhandler
## task = ["STORE", key, origfw, curfw, reason, TaskComplRoutine, \
##         "na",] # timestamp to be appended by complhandler
    def _parse_task(self, task):
        taskType = task[TSK_TASKTYPE]
        taskCwid = task[TSK_CWID]
        taskMac = task[TSK_MAC]
        taskOrig = task[TSK_ORIG]
        taskCurr = task[TSK_CURR]
        taskStat = task[TSK_STAT]
        taskFwType = task[TSK_FWTYPE]
        taskTs = task[TSK_TIME]
        return (taskType, taskCwid, taskMac, taskOrig, taskCurr,\
                taskStat, taskFwType, taskTs)

    def start_log(self):
        try:
            if (os.path.exists(self.logfilepath)):    
                self.filehandle = open(self.logfilepath, 'a')
            else:
                self.filehandle = open(self.logfilepath, 'w')
                self.filehandle.write("\n(started %s)\n" % datetime.datetime.now())
                self.filehandle.write("====================================\n")
        except Exception, ex:
            print "_start_log: exception= ", ex
            return False
        else:
            return True

    def log_task(self, task):
        try:
            print "log_task: enter"
            origfw = "NA"
            curfw = "NA"
            taskinfo = self._parse_task(task)
            tasktype, cwid, mac, origfw, curfw, status, fwType, ts = taskinfo
            print "logtask: ", taskinfo
            getch()
            
            print "log_task: taskinfo= ", taskinfo 
            timestampLine = "[" + ts + "]" + "\n"
            print timestampLine.strip()
            deviceLine = "Device: " + \
                         cwid + " (" + mac + ")" + "\n"
            print deviceLine.strip()
            firmwareLine = "Original: " + origfw +"," + " Current: " + \
                           curfw + "\n"
            print firmwareLine.strip()
            taskLine = "Task processed: " + tasktype + "\n"
            statusLine = "Reason: " + status + "\n\n"
            print statusLine
            self.filehandle.write("Device processed on: ")
            self.filehandle.write(timestampLine)
            self.filehandle.write(deviceLine)
            self.filehandle.write(firmwareLine)
            self.filehandle.write(taskLine)
            self.filehandle.write(statusLine)
            print "log_task: closing log file..."
            self.filehandle.close()
        except Exception, ex:
            print "log_task: exception= ", ex
            self.filehandle.close()
            return False
        else:
            print "log_task: exit"
            return True        

def TaskComplHandler(completedTask, logfilepath):
    try:
        logger = Logger(logfilepath)
        ret = logger.start_log()
        if (ret == False):
            print "TaskComplHandler: "
            print " ## WARNING: NOT LOGGING TO FILE BECAUSE OF "
            print " POSSIBLE FILE IO ERROR ##"
        else:
            print "\n*** TaskCompleteHandler ***"
            timestamp = strftime("%d-%m-%Y %H:%M:%S")
            completedTask.append(timestamp)
            print "*** Task type = ", completedTask[TSK_TASKTYPE]
            print "*** Device = ", completedTask[TSK_CWID]
            print "*** Device MAC = ", completedTask[TSK_MAC]
            print "*** Original firmware = ", completedTask[TSK_ORIG]
            print "*** Updated firmware = ", completedTask[TSK_CURR]
            print "*** Update status = ", completedTask[TSK_STAT]
            print "*** Time & Date stamp = ", completedTask[TSK_TIME]
            print "*** Logging to file..."        
            ret = logger.log_task(completedTask)
    except Exception, ex:
        print "TaskComplHandler: exception= ", ex
        ret = False
    else:
        return ret
    
class XBeeFwUpgrade:

    def __init__(self, idigiInfo, fwInfo, logfilepath, interval=60000):
        self._idigiInfo = idigiInfo
        self.fwFileName = fwInfo[1]
        self.requiredver_zb = fwInfo[2]
        self._encodedFwImg = ""
        self._taskQueue = Queue(0)
        self._stopMonitoring = True
        self.iDigiCom = idigi_comm.WebServiceComm(idigiInfo)
        self.processing = False
        self.updateInterval = 60000 # defaults to 10min
        self.logFilePath = logfilepath #= Logger(logfilepath)
        # master list, holds all devices obtained from iDigi
        self.idigiDevices = IDigiDeviceDict()
        self.skiplist = []
        self.hw_type_ok = []
        self.retries = 3

    def blast_upgrade(self):
        try:
            retString = self.iDigiCom._get_cwid_list_from_device_core(self)
            print "retString = ", retString
        except:
            return False
        else:
            return True

    def _encode_store_fw(self, fwpath):
        try:
            # DO: allow for ALL upgrades!
            # prepare firmware file to store encoded image
            hFile = open(fwpath, "rb")
            fileData = hFile.read()
            hFile.close()
            tmpImg = base64.encodestring(fileData)
            print "len of tmpImg= ", len(tmpImg)
            self._encodedFwImg = tmpImg[:len(tmpImg)-1]
            print "len of encImg= ", len(self._encodedFwImg)

        except IOError:
            print "\nWARNING: Cannot open firmware file at location specified."
            print "Please check the config file and verify the "
            print "correct firmware image path is listed. Firmware "
            print "files will end in an EBL extension.\n"
            print "(PRESS RETURN TO EXIT UPGRADE SCRIPT)"
            getch()
            sys.exit(0)
        except Exception, ex:
            print "prep_firmware: exception= ", ex
            return "ERROR"
        else:
            return True

    # returns (status, descr if error)
    def _upgrade_fw(self, cwid):
        try:
            print "_upgrade_fw: enter: ", cwid
            retStatus = False
            xmlreply = ""
            # send image
            print "_uprgrade_fw: sending image"
            xmlrepl = self.iDigiCom.send_rci_request(cwid, \
                            "UploadFw", self.fwFileName, \
                                self._encodedFwImg)
            if (self.iDigiCom.lastStatusCode == 200):
                print "_upgrade_fw: checking reply"
                
                errmsg = idigi_parse.check_for_errors(xmlrepl)
                
                if (errmsg != "OK"):
                    print "_upgrade_fw: returning retstatus"
                    retStatus = ("ERROR", errmsg)
                else:
                    xmlreply = ""
                    print "send upd cmd"
                    xmlreply = self.iDigiCom.send_rci_request(cwid, \
                                    "SendUpdateCmd", self.fwFileName)
                    # check for last send status
                    if (self.iDigiCom.lastStatusCode == 200):
                        print "_upgrade_fw: checking reply"
                        errmsg = idigi_parse.check_for_errors(xmlrepl)

                        if (errmsg != "OK"):
                            retStatus = ("ERROR", errmsg)
                        else:
                            retStatus = ("OK", "SUCCESS")
                    else:
                        retStatus = ("ERROR", "Error sending upg command")
            else:
                retStatus = ("ERROR", "Error sending firmware image")
        except Exception, ex:
            print "_upgrade_fw: exception= ", ex
        else:
            print "returning retstatus ", retStatus
            return retStatus

    def _delete_fw_image(self, cwid):
        try:
            print "_delete_fw_image: enter"
            retStatus = self.iDigiCom.send_rci_request(cwid, \
                            "DeleteFile", self.fwFileName)
            print "_delete_fw_image: status = ", retStatus
        except Exception, ex:
            return "_delete_fw_image: exception= ", ex
        else:
            return retStatus

    # used in original dictionary build (cached = True)
    # used to verify firmware upgrade (cached = False)
    def _get_current_ver_info(self, targetList, refresh, printXml=True):
        try:
            print "_get_current_ver_info:"
            resp = self.iDigiCom.query_state(targetList, refresh, \
                                             "<zigbee_state/>")
            if (resp != "ERROR"):
                parsed = idigi_parse.parse_query_state_resp(resp)
            if (printXml):
                print resp
                print parsed
        except Exception, ex:
            print "_get_current_ver_info: exception= ", ex
            return "ERROR"
        else:
            return parsed

    def upgrade_dev_fw(self, cwid):
        try:
            print "upgrade_dev_fw: enter"
            upgStat = []
            verified = False
            # send fw image and then upgr command            
            retStatus = self._upgrade_fw(cwid)

            if (retStatus[0] == "ERROR"):
                print "upgrade_dev_fw: returning status ", retStatus
                return retStatus[0], retStatus[1], verified
                #return retStatus, verified
            #else:
            #    return ("SUCCESS", "Device updated successfully")

            #verifyEnabled = True

            # verification enabled?
            #if (verifyEnabled == True):
            # verify fw upgrade
            # give us the uncached idigi response...
            print "upgrade_dev_fw: VERIFYING..."
            ret = self._get_current_ver_info(cwid, True)
            print "upgrade_dev_fw: RET= ", ret["firmware_version"]
            print "upgrade_dev_fw: REQ= ", self.requiredver_zb

            if (ret["firmware_version"] == self.requiredver_zb):
                print "upgrade_dev_fw: upgrade verified"
                verified = True
            else:
                print "upgrade_dev_fw: upgrade verification failed"
                verified = False
            # TODO:       
            # delete file enabled?
            #if (deleteFwFile == True):
            #    status = self._delete_fw_image(cwid)
            #return status

        except Exception, ex:
            print "upgrade_dev_fw: an error has occured = ", ex
            return "ERROR", "Application error", False
        else:
            #return retStatus, verified
            print "upgrade_dev_fw: returning= ",
            print retStatus[0]
            print retStatus[1]
            print verified 
            return retStatus[0], retStatus[1], verified

    # queue related    
    def enqueue(self, task):
        try:
            print "TASK QUEUED: "
            print task
            self._taskQueue.put(task)
            return True
        except Exception, ex:
            print "enqueue: exception: ", ex
            return "Error"

    def dequeue(self):
        try:
            if (self._taskQueue.empty()):
                return (None, None, None, None, None, None)
            else:
                print "dequeue: Taking off a task!"
                return self._taskQueue.get()
        except Exception, ex:
            print "dequeue: exception: ", ex
            return "Error"

    def _get_compare_data(self, device):
        try:
            print "_get_compare_data: enter"
            conn = "na"
            curhw = "na"
            curfw = "na"
            hwtype = "na"
            fwtype = "na"
            assoc = "na"
            print "_get_compare_data"
            keylist = device.keys()
            if "dpConnectionStatus" in keylist:
                conn = device["dpConnectionStatus"]
            if "hardware_version" in keylist:
                curhw = device["hardware_version"]
            if "firmware_version" in keylist:
                curfw = device["firmware_version"]
            if "association" in keylist:
                assoc = device["association"]
            if (curhw != "na") and (curfw != "na"):
                hwtype = curhw[0:4]
                print "HWTYPE= ", hwtype
                fwtype = curfw[2:3]

            compareData = (conn, curhw, curfw, hwtype, fwtype, assoc)

        except Exception, ex:
            return "ERROR= ", ex
        else:
            return compareData

    #if not online, queues "STORE" task. If online, sends to __process
    def process_devices(self, ProcComplHandler=None, specialFlags=None):
        try:
            upgCnt = 0
            noChange = 0
            upgTaskCount = 0
            noChangeCount = 0
            print "**process_devices**"

            for key in self.idigiDevices.keys():
                devMac = self.idigiDevices[key]["devMac"]
                if  (self.idigiDevices[key]["dpConnectionStatus"] == "0"):
                    # don't even bother!
                    task = ["STORE", key, devMac, "NA", "NA", "Not Online", \
                            TaskComplHandler, "NA"]
                elif (self.idigiDevices[key]["dpConnectionStatus"] == "1"):
                    (upgCnt, noChange, task) = self.__process(key, \
                                                    self.idigiDevices[key], \
                                                    devMac, TaskComplHandler,\
                                                    specialFlags)
                upgTaskCount += upgCnt
                noChangeCount += noChange
                self.enqueue(task)
            # queue up a quit task for the thread to quit...
            self.enqueue(["QUIT"])
        except Exception, ex:
            print "process_devices: exception = ", ex
        else:
            return (upgTaskCount, noChangeCount)

    # processes each individual connectware device
    def __process(self, key, device, devMac, TaskComplRoutine, specialFlags=None,fwTypeList=None):
        try:
            # Note: pre-req to get here is being 'connected' to idigi
            upgradeNeeded = False
            upgrCount = 0
            skipCount = 0
            reason = "na"
            task = []

            compData = self._get_compare_data(device)
            if ("ERROR" or "na") not in compData:
                conn, curhw, curfw, hwtype, fwtype, assoc = compData
            else:
                upgradeNeeded = False
                reason = "Could not convert current device data"

            zbfw_required = self.requiredver_zb
            print "__process: zb fw required = ", zbfw_required
            print "__process: special Flags = ", specialFlags
            
            # check radio type first
            radiotype = radio_types[hwtype]

            if (radiotype in self.hw_type_ok):
                # check firmware next
                if (fwtype == ZB_FW):
                    # last three digits can be conv to ints & compared 
                    intCurVer = int(curfw[3:6], 10)
                    intReqVer = int(zbfw_required[3:6], 10)
                    if (specialFlags['forceupg'] == True):
                        print "forcing upgrade"
                        upgradeNeeded = True
                        reason = "ZB: force upgrade flag was set to True"
                        upgrCount = 1
                    elif (intCurVer < intReqVer):
                        print "upgrade needed"
                        upgradeNeeded = True
                        reason = "ZB: old firmware version"
                        upgrCount = 1
                    else:
                        upgradeNeeded = False
                        reason = "ZB: firmware is up to date" 
                        skipCount += 1
                elif (fwtype == ZNET_FW):
                    if  (assoc == "0x0"):
                        upgradeNeeded = False
                        reason = "ZNet: associated"
                        skipCount += 1
                    elif (assoc != "0x0"):
                        upgradeNeeded = True
                        reason = "ZNet: not associated"
                        upgrCount = 1
                    else:
                        upgradeNeeded = False
                        reason = "ZNet: cannot determine association"
                        skipCount += 1
                else:
                    upgradeNeeded = False
                    reason = "Non-ZB or Non-ZNet firmware"
                    skipCount += 1
            else:
                upgradeNeeded = False
                reason = "Incorrect hardware type"
                skipCount += 1

            if (upgradeNeeded == True):
                ## we have a cached vers but also know its one we want to
                ## upgrade at this point, hit with uncached query (give us
                ## definite.  If device is running current version, do the
                ## store, otherwise do the upgrade
                task = ["UPGR", key, devMac, curfw, curfw, reason, \
                        TaskComplRoutine, \
                        ZB_FW] # +timestamp to be appended by complhandler
            else:
                task = ["STORE", key, devMac, curfw, curfw, reason, \
                        TaskComplRoutine, \
                        "na"] # +timestamp to be appended by complhandler

        except Exception, ex:
            print "__process: exception= ", ex
            return (0, task)
        else:
            return (upgrCount, skipCount, task)

    def add_radio_info_to_devices(self, printToScr=False):
        try:
            for key in self.idigiDevices.keys():
                print key
                curHw = ""
                curFW = ""
                if (self.idigiDevices[key]["dpConnectionStatus"] == "1"):
                    # give us cached idigi response for initial dict build
                    zbsDict = self._get_current_ver_info(key, refresh=False)
                    print "\n\nzbs=", zbsDict
                    # take out entry with duplicate cwid
                    del zbsDict["device"]
                    print "add_radio_info_to_devices"
                    self.idigiDevices[key].update(zbsDict)
        except Exception, ex:
            print "add_radio_info_to_devices: exception= ", ex
        else:
            return True

    # uses idigi_parse.parse_devcore_resp to build device dictionary
    def build_device_dict(self, printToScr=False):
        try:
            self.idigiDevices = {}
            resp = self.iDigiCom.send_msg("/ws/DeviceCore", "GET")

            if (printToScr):
                print resp            
                # DO: retParsed = self.iDigiCom.get_devices(resp)
                #                 which calls build_idigi_devices
                #               which was once "parse_for_devices"
            self.idigiDevices = idigi_parse.build_devices_from_xml(resp)

        except Exception, ex:
            print "build_device_dict: exception= ", ex
            return "ERROR"
        else:
            return True                    

    def set_firmware_file(self, fullpath):
        retStatus = self._encode_store_fw(fullpath)
        return retStatus

    def set_hwtype_ok(self, hw_type_ok_list):
        self.hw_type_ok = hw_type_ok_list
        return True
