#!/usr/bin/env python

import zerorpc
from jnpr.junos import Device
from jnpr.junos.utils.config import Config
from jnpr.junos.exception import *
from jxmlease import parse
from lxml import etree

class ZPyEZ(object):
    _devModel = {'host': '', 'user':'', 'password':''}
    _devices = []
    _devicesSet = []
    _configs = []
    
    def checkInput(self, _dev):
        # Basic input check to verify that all keys (host, user, password) exist 
        if len(set(self._devModel.keys()).difference(_dev.keys())) > 0:
                raise Exception("host, user and password Keys are needed." )
        # To add regex verification of the Values as a next step as _devModel Values
        # Modify function to accept other PyEZ Dev Keys/Values inputs 
        pass 
    
    
    def addDevice(self, devToAdd):
        if devToAdd.keys()[0] in set(self._devicesSet):
            index = 0 
            for device in self._devices:
                if device.keys()[0] == devToAdd.keys()[0]:
                    self._devices.pop(index)
                    self._devices.append(devToAdd)
                else: 
                    index += 1
        else:
            self._devices.append(devToAdd)
            self._devicesSet.append(devToAdd.keys()[0])
        pass  
    
    def addConfig(self, devToAdd):
        self._configs.append(devToAdd)
    
    def getDevice(self, devId): 
        if devId in set(self._devicesSet): 
            for device in self._devices: 
                if device.keys()[0] == devId:
                    return device.values()[0]
        else: 
            return None

    def _getConfig(self, devId): 
        for config in self._configs: 
            if config.keys()[0] == devId: 
                return config.values()[0] 
        return None
        
    def delDevice(self, devId): 
        if devId in set(self._devicesSet):
            index = 0 
            for device in self._devices: 
                if device.keys()[0] == devId:
                    try: 
                        rmDev = device.values()[0]
                        rmDev.close()
                        self._devices.pop(index)
                        index2=0
                        for _devId in self._devicesSet:
                            if _devId == devId:
                                self._devicesSet.pop(index2)
                            else: 
                                index2 += 1
                        index3=0 
                        for config in self._configs: 
                            if config.keys()[0] == devId:
                                self._configs.pop(index3)
                            else: 
                                index3 += 1 
                        break 
                    except:
                        self._devices.pop(index)
                        index2=0
                        for _devId in self._devicesSet:
                            if _devId == devId:
                                self._devicesSet.pop(index2)
                            else: 
                                index2 += 1
                        break
                else: 
                    index += 1
            
        else: 
            return None
        
    def open(self, _dev):
        """
            Opening Device with input matching {'host': '', 'user'='', 'password'=''}
        """
        self.checkInput(_dev)
        try: 
            dev = Device (**_dev)
            dev.open()
            
            # Adding PyEZ Object reference to an Array - keyed by devId  
            self.addDevice({str(dev): dev})
            
            # returning devId 
            return str(dev) 
        except ConnectError as err:
            raise  Exception("Cannot connect to device: {0}".format(err)) 
            
    def config(self, _devId): 
        dev = self.getDevice(_devId)
        if dev != None:
            try:
                cu = Config(dev)
                self.addConfig({str(dev): cu})
                
                # returning devId 
                return str(dev) 
            except: 
                raise  Exception("Cannot Configure Device: {0}".format(_devId))
        else: 
            raise  Exception("Object cannot be found: {0}".format(_devId))
            
    def load(self, _devId, _configStatement):
        cu = self._getConfig(_devId)
        if cu != None: 
            cu.load(_configStatement)
        else:
            raise  Exception("Config cannot be found: {0}".format(_devId))
    
    def load_text(self, _devId, _configText, _merge=True):
        cu = self._getConfig(_devId)
        if cu != None: 
            cu.load(_configText, format="text", merge=_merge)
        else:
            raise  Exception("Config cannot be found: {0}".format(_devId))      
        
    def pdiff(self, _devId):
        cu = self._getConfig(_devId)
        if cu != None: 
            return cu.diff()
        else:
            raise  Exception("Config cannot be found: {0}".format(_devId))
    
    def commit_check(self, _devId):
        cu = self._getConfig(_devId)
        if cu != None: 
            return cu.commit_check()
        else:
            raise  Exception("Config cannot be found: {0}".format(_devId))
    
    def commit(self, _devId):
        cu = self._getConfig(_devId)
        if cu != None: 
            return cu.commit()
        else:
            raise  Exception("Config cannot be found: {0}".format(_devId))
    
    def rollback(self, _devId, _rollbackId=0):
        cu = self._getConfig(_devId)
        if cu != None: 
            return cu.rollback(_rollbackId)
        else:
            raise  Exception("Config cannot be found: {0}".format(_devId))
        
        
    def facts(self, _devId):
        dev = self.getDevice(_devId)
        if dev != None:
            r = dev.facts 
            del r['version_info']
            return r
        else: 
            raise  Exception("Object cannot be found: {0}".format(_devId))
    
    def hostname(self, _devId):
        dev = self.getDevice(_devId)
        if dev != None: 
            return dev.hostname 
        else: 
            raise  Exception("Object cannot be found: {0}".format(_devId))
    
    def connected(self, _devId):
        dev = self.getDevice(_devId)
        if dev != None: 
            return dev.connected 
        else: 
            raise  Exception("Object cannot be found: {0}".format(_devId))
            
    def facts_refresh(self, _devId):
        dev = self.getDevice(_devId)
        if dev != None: 
            return dev.facts_refresh() 
        else: 
            raise  Exception("Object cannot be found: {0}".format(_devId))
    
    def ON_JUNOS(self, _devId):
        dev = self.getDevice(_devId)
        if dev != None: 
            return dev.ON_JUNOS 
        else: 
            raise  Exception("Object cannot be found: {0}".format(_devId))
    
    def port(self, _devId):
        dev = self.getDevice(_devId)
        if dev != None: 
            return dev.port 
        else: 
            raise  Exception("Object cannot be found: {0}".format(_devId))
    
    def manages(self):
        return self._devicesSet
    
    def manages_dev(self):
        retDevs = []
        for dev in self._devices:
            retDevs.append(dev.keys()[0])
        return retDevs
    
    def close_device(self, _devId):
        dev = self.getDevice(_devId)
        if dev != None: 
            self.delDevice(_devId) 
        else: 
            raise  Exception("Object cannot be found: {0}".format(_devId))
    
    def get_config(self, _devId, format="xml"):
        dev = self.getDevice(_devId)
        if dev != None: 
            if format == "xml":
                _resp = dev.rpc.get_config()
                return  etree.tostring(_resp).strip()
            elif format == "text":
                _resp = dev.rpc.get_configuration(dict(format="text"))
                return _resp.text 
            elif format =="set":
                _resp = dev.cli('show configuration | display set', warning=False)
                return _resp 
            else: 
                raise  Exception("format not Supported: {0}".format(format))  
        else: 
            raise  Exception("Object cannot be found: {0}".format(_devId))
    
    def display_xml_rpc(self, _devId, devCommand):
        dev = self.getDevice(_devId)
        if dev != None: 
            return etree.tostring(dev.display_xml_rpc(devCommand)).strip() 
        else: 
            raise  Exception("Object cannot be found: {0}".format(_devId))
        
    def rpc(self, _devId, cliCommand):
        dev = self.getDevice(_devId)
        if dev != None: 
            xml_rpc = dev.display_xml_rpc(cliCommand)
            rpcMethod = getattr(dev, 'rpc')
            try: 
                _resp= rpcMethod(xml_rpc)
                return etree.tostring(_resp).strip()
            except: 
                raise  Exception("Error while executing : {0}".format(cliCommand))
        else: 
            raise  Exception("Object cannot be found: {0}".format(_devId))
            
        
    def xml_rpc(self, _devId, xmlRPC):
        dev = self.getDevice(_devId)
        if dev != None: 
            rpcMethod = getattr(dev, 'rpc')
            try: 
                _resp= rpcMethod(xmlRPC)
                return etree.tostring(_resp).strip()
            except: 
                raise  Exception("Error while executing : {0}".format(xmlRPC))
        else: 
            raise  Exception("Object cannot be found: {0}".format(_devId))
        
    def cli(self, _devId, cliCommand):
        dev = self.getDevice(_devId)
        if dev != None: 
            try: 
                _ret = dev.cli(cliCommand, warning=False)
                return _ret 
            except:
                raise Exception("Error while executing command: {0}".format(cliCommand)) 
        else: 
            raise  Exception("Object cannot be found: {0}".format(_devId))



## main 
if __name__ == "__main__":
    s = zerorpc.Server(ZPyEZ())
    s.bind("tcp://0.0.0.0:8484")
    s.run()