drivers/CiscoSX300.py

#! /usr/bin/python

#  Copyright 2014 Linaro Limited
#
#  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., 51 Franklin Street, Fifth Floor, Boston,
#  MA 02110-1301, USA.

import logging
import pexpect
import sys
import re

if __name__ == '__main__':
    import os
    vlandpath = os.path.abspath(os.path.normpath(os.path.dirname(sys.argv[0])))
    sys.path.insert(0, vlandpath)
    sys.path.insert(0, "%s/.." % vlandpath)

from errors import InputError, PExpectError
from drivers.common import SwitchDriver, SwitchErrors

class CiscoSX300(SwitchDriver):

    connection = None
    _username = None
    _password = None
    _enable_password = None

    # No extra capabilities for this switch/driver yet
    _capabilities = [
    ]

    # Regexp of expected hardware information - fail if we don't see
    # this
    _expected_descr_re = re.compile('.*\d+-Port.*Managed Switch.*')

    def __init__(self, switch_hostname, switch_telnetport=23, debug = False):
        SwitchDriver.__init__(self, switch_hostname, debug)
        self._systemdata = []
        self.exec_string = "/usr/bin/telnet %s %d" % (switch_hostname, switch_telnetport)
        self.errors = SwitchErrors()

    ################################
    ### Switch-level API functions
    ################################

    # Connect to the switch and log in
    def switch_connect(self, username, password, enablepassword):
        self._username = username
        self._password = password
        self._enable_password = enablepassword
        self._switch_connect()

    # Log out of the switch and drop the connection and all state
    def switch_disconnect(self):
        self._logout()
        logging.debug("Closing connection: %s", self.connection)
        self.connection.close(True)
        self._ports = []
        self._prompt_name = ''
        self.systemdata = []
        del(self)

    # Save the current running config into flash - we want config to
    # remain across reboots
    def switch_save_running_config(self):
        try:
            self._cli("copy running-config startup-config")
            self.connection.expect("Y/N")
            self._cli("y")
            self.connection.expect("succeeded")
        except (PExpectError, pexpect.EOF, pexpect.TIMEOUT):
            # recurse on error
            self._switch_connect()
            self.switch_save_running_config()

    # Restart the switch - we need to reload config to do a
    # roll-back. Do NOT save running-config first if the switch asks -
    # we're trying to dump recent changes, not save them.
    #
    # This will also implicitly cause a connection to be closed 
    def switch_restart(self):
        self._cli("reload")
        index = self.connection.expect(['Are you sure', 'will reset'])
        if index == 0:
            self._cli("y") # Yes, continue without saving
            self.connection.expect("reset the whole")

        # Fall through
        self._cli("y") # Yes, continue to reset
        self.connection.close(True)

    ################################
    ### VLAN API functions
    ################################

    # Create a VLAN with the specified tag
    def vlan_create(self, tag):
        logging.debug("Creating VLAN %d", tag)

        try:
            self._configure()
            self._cli("vlan database")
            self._cli("vlan %d" % tag)
            self._end_configure()

            # Validate it happened
            vlans = self.vlan_get_list()
            for vlan in vlans:
                if vlan == tag:
                    return
            raise IOError("Failed to create VLAN %d" % tag)

        except PExpectError:
            # recurse on error
            self._switch_connect()
            self.vlan_create(tag)

    # Destroy a VLAN with the specified tag
    def vlan_destroy(self, tag):
        logging.debug("Destroying VLAN %d", tag)

        try:
            self._configure()
            self._cli("no vlan %d" % tag)
            self._end_configure()

            # Validate it happened
            vlans = self.vlan_get_list()
            for vlan in vlans:
                if vlan == tag:
                    raise IOError("Failed to destroy VLAN %d" % tag)

        except PExpectError:
            # recurse on error
            self._switch_connect()
            self.vlan_destroy(tag)

    # Set the name of a VLAN
    def vlan_set_name(self, tag, name):
        logging.debug("Setting name of VLAN %d to %s", tag, name)

        try:
            self._configure()
            self._cli("vlan %d" % tag)
            self._cli("interface vlan %d" % tag)
            self._cli("name %s" % name)
            self._end_configure()

            # Validate it happened
            read_name = self.vlan_get_name(tag)
            if read_name != name:
                raise IOError("Failed to set name for VLAN %d (name found is \"%s\", not \"%s\")"
                              % (tag, read_name, name))

        except PExpectError:
            # recurse on error
            self._switch_connect()
            self.vlan_set_name(tag, name)

    # Get a list of the VLAN tags currently registered on the switch
    def vlan_get_list(self):
        logging.debug("Grabbing list of VLANs")

        try:
            vlans = []
            regex = re.compile('^ *(\d+).*(D|S|G|R)')

            self._cli("show vlan")
            for line in self._read_long_output("show vlan"):
                match = regex.match(line)
                if match:
                    vlans.append(int(match.group(1)))
            return vlans

        except PExpectError:
            # recurse on error
            self._switch_connect()
            return self.vlan_get_list()

    # For a given VLAN tag, ask the switch what the associated name is
    def vlan_get_name(self, tag):
        logging.debug("Grabbing the name of VLAN %d", tag)

        try:
            name = None
            regex = re.compile('^ *\d+\s+(\S+).*(D|S|G|R)')
            self._cli("show vlan tag %d" % tag)
            for line in self._read_long_output("show vlan tag"):
                match = regex.match(line)
                if match:
                    name = match.group(1)
            name.strip()
            return name

        except PExpectError:
            # recurse on error
            self._switch_connect()
            return self.vlan_get_name(tag)

    ################################
    ### Port API functions
    ################################    

    # Set the mode of a port: access or trunk
    def port_set_mode(self, port, mode):
        logging.debug("Setting port %s to %s", port, mode)
        if not self._is_port_mode_valid(mode):
            raise InputError("Port mode %s is not allowed" % mode)
        if not self._is_port_name_valid(port):
            raise InputError("Port name %s not recognised" % port)

        try:
            self._configure()
            self._cli("interface %s" % port)
            self._cli("switchport mode %s" % mode)
            self._end_configure()

            # Validate it happened
            read_mode = self.port_get_mode(port)
            if read_mode != mode:
                raise IOError("Failed to set mode for port %s" % port)

        except PExpectError:
            # recurse on error
            self._switch_connect()
            self.port_set_mode(port, mode)

    # Get the mode of a port: access or trunk
    def port_get_mode(self, port):
        logging.debug("Getting mode of port %s", port)
        mode = ''
        if not self._is_port_name_valid(port):
            raise InputError("Port name %s not recognised" % port)
        regex = re.compile('Port Mode: (\S+)')

        try:
            self._cli("show interfaces switchport %s" % port)
            for line in self._read_long_output("show interfaces switchport"):
                match = regex.match(line)
                if match:
                    mode = match.group(1)
            return mode.lower()

        except PExpectError:
            # recurse on error
            self._switch_connect()
            return self.port_get_mode(port)

    # Set an access port to be in a specified VLAN (tag)
    def port_set_access_vlan(self, port, tag):
        logging.debug("Setting access port %s to VLAN %d", port, tag)
        if not self._is_port_name_valid(port):
            raise InputError("Port name %s not recognised" % port)
        if not (self.port_get_mode(port) == "access"):
            raise InputError("Port %s not in access mode" % port)

        try:
            self._configure()
            self._cli("interface %s" % port)
            self._cli("switchport access vlan %d" % tag)
            self._end_configure()

            # Validate things worked
            read_vlan = int(self.port_get_access_vlan(port))
            if read_vlan != tag:
                raise IOError("Failed to move access port %s to VLAN %d - got VLAN %d instead"
                              % (port, tag, read_vlan))

        except PExpectError:
            # recurse on error
            self._switch_connect()
            self.port_set_access_vlan(port, tag)

    # Add a trunk port to a specified VLAN (tag)
    def port_add_trunk_to_vlan(self, port, tag):
        logging.debug("Adding trunk port %s to VLAN %d", port, tag)
        if not self._is_port_name_valid(port):
            raise InputError("Port name %s not recognised" % port)
        if not (self.port_get_mode(port) == "trunk"):
            raise InputError("Port %s not in trunk mode" % port)

        try:
            self._configure()
            self._cli("interface %s" % port)
            self._cli("switchport trunk allowed vlan add %d" % tag)
            self._end_configure()

            # Validate it happened
            read_vlans = self.port_get_trunk_vlan_list(port)
            for vlan in read_vlans:
                if vlan == tag:
                    return
            raise IOError("Failed to add trunk port %s to VLAN %d" % (port, tag))

        except PExpectError:
            # recurse on error
            self._switch_connect()
            self.port_add_trunk_to_vlan(port, tag)
        
    # Remove a trunk port from a specified VLAN (tag)
    def port_remove_trunk_from_vlan(self, port, tag):
        logging.debug("Removing trunk port %s from VLAN %d", port, tag)
        if not self._is_port_name_valid(port):
            raise InputError("Port name %s not recognised" % port)
        if not (self.port_get_mode(port) == "trunk"):
            raise InputError("Port %s not in trunk mode" % port)

        try:
            self._configure()
            self._cli("interface %s" % port)
            self._cli("switchport trunk allowed vlan remove %d" % tag)
            self._end_configure()

            # Validate it happened
            read_vlans = self.port_get_trunk_vlan_list(port)
            for vlan in read_vlans:
                if vlan == tag:
                    raise IOError("Failed to remove trunk port %s from VLAN %d" % (port, tag))

        except PExpectError:
            # recurse on error
            self._switch_connect()
            self.port_remove_trunk_from_vlan(port, tag)

    # Get the configured VLAN tag for an access port (tag)
    def port_get_access_vlan(self, port):
        logging.debug("Getting VLAN for access port %s", port)
        vlan = 1
        if not self._is_port_name_valid(port):
            raise InputError("Port name %s not recognised" % port)
        if not (self.port_get_mode(port) == "access"):
            raise InputError("Port %s not in access mode" % port)
        regex = re.compile('(\d+)\s+\S+\s+Untagged\s+(Static|System)')

        try:
            self._cli("show interfaces switchport %s" % port)
            for line in self._read_long_output("show interfaces switchport"):
                match = regex.match(line)
                if match:
                    vlan = match.group(1)
            return int(vlan)

        except PExpectError:
            # recurse on error
            self._switch_connect()
            return self.port_get_access_vlan(port)

    # Get the list of configured VLAN tags for a trunk port
    def port_get_trunk_vlan_list(self, port):
        logging.debug("Getting VLANs for trunk port %s", port)
        vlans = [ ]
        if not self._is_port_name_valid(port):
            raise InputError("Port name %s not recognised" % port)
        if not (self.port_get_mode(port) == "trunk"):
            raise InputError("Port %s not in trunk mode" % port)
        regex = re.compile('(\d+)\s+\S+\s+(Tagged|Untagged)\s+Static')

        try:
            self._cli("show interfaces switchport %s" % port)
            for line in self._read_long_output("show interfaces switchport"):
                match = regex.match(line)
                if match:
                    vlans.append (int(match.group(1)))
            return vlans

        except PExpectError:
            # recurse on error
            self._switch_connect()
            return self.port_get_trunk_vlan_list(port)

    ################################
    ### Internal functions
    ################################

    # Connect to the switch and log in
    def _switch_connect(self):

        if not self.connection is None:
            self.connection.close(True)
            self.connection = None

        logging.debug("Connecting to Switch with: %s", self.exec_string)
        self.connection = pexpect.spawn(self.exec_string, logfile = self.logfile)

        self._login()

        # Avoid paged output
        self._cli("terminal datadump")

        # And grab details about the switch. in case we need it
        self._get_systemdata()

        # And also validate them - make sure we're driving a switch of
        # the correct model! Also store the serial number
        descr_regex = re.compile('System Description:.\s+(.*)')
        sn_regex = re.compile('SN:\s+(\S_)')
        descr = ""

        for line in self.systemdata:
            match = descr_regex.match(line)
            if match:
                descr = match.group(1)
            match = sn_regex.match(line)
            if match:
                self.serial_number = match.group(1)

        if not self._expected_descr_re.match(descr):
            raise IOError("Switch %s not recognised by this driver: abort" % descr)

        # Now build a list of our ports, for later sanity checking
        self._ports = self._get_port_names()
        if len(self._ports) < 4:
            raise IOError("Not enough ports detected - problem!")

    def _login(self):
        logging.debug("attempting login with username %s, password %s", self._username, self._password)
        self._cli("")
        self.connection.expect("User Name:")
        self._cli("%s" % self._username)
        self.connection.expect("Password:")
        self._cli("%s" % self._password, False)
        self.connection.expect("\*\*")
        while True:
            index = self.connection.expect(['User Name:', 'authentication failed', r'([^#]+)#', 'Password:', '.+'])
            if index == 0 or index == 1: # Failed to log in!
                logging.error("Login failure: %s\n", self.connection.match)
                raise IOError
            elif index == 2:
                self._prompt_name = re.escape(self.connection.match.group(1).strip())
                # Horrible output from the switch at login time may
                # confuse our pexpect stuff here. If we've somehow got
                # multiple lines of output, clean up and just take the
                # *last* line here. Anything before that is going to
                # just be noise from the "***" password input, etc.
                prompt_lines = self._prompt_name.split('\r\n')
                if len(prompt_lines) > 1:
                    self._prompt_name = prompt_lines[-1]
                logging.debug("Got prompt name %s", self._prompt_name)
                return 0
            elif index == 3 or index == 4:
                self._cli("", False)

    def _logout(self):
        logging.debug("Logging out")
        self._cli("exit", False)

    def _configure(self):
        self._cli("configure terminal")

    def _end_configure(self):
        self._cli("end")

    def _read_long_output(self, text):
        prompt = self._prompt_name + '#'
        try:
            self.connection.expect(prompt)
        except (pexpect.EOF, pexpect.TIMEOUT):
            # Something went wrong; logout, log in and try again!
            logging.error("PEXPECT FAILURE, RECONNECT")
            self.errors.log_error_in(text)
            raise PExpectError("_read_long_output failed")
        except:
            logging.error("prompt is \"%s\"", prompt)
            raise

        longbuf = []
        for line in self.connection.before.split('\r\n'):
            longbuf.append(line.strip())
        return longbuf

    def _get_port_names(self):
        logging.debug("Grabbing list of ports")
        interfaces = []

        # Use "Up" or "Down" to only identify lines in the output that
        # match interfaces that exist
        regex = re.compile('^(\w+).*(Up|Down)')

        try:
            self._cli("show interfaces status detailed")
            for line in self._read_long_output("show interfaces status detailed"):
                match = regex.match(line)
                if match:
                    interfaces.append(match.group(1))
            logging.debug("  found %d ports on the switch", len(interfaces))
            return interfaces
        except PExpectError:
            # recurse on error
            self._switch_connect()
            return self._get_port_names()

    def _show_config(self):
        logging.debug("Grabbing config")
        try:
            self._cli("show running-config")
            return self._read_long_output("show running-config")
        except PExpectError:
            # recurse on error
            self._switch_connect()
            return self._show_config()

    def _show_clock(self):
        logging.debug("Grabbing time")
        try:
            self._cli("show clock")
            return self._read_long_output("show clock")
        except PExpectError:
            # recurse on error
            self._switch_connect()
            return self._show_clock()

    def _get_systemdata(self):
        logging.debug("Grabbing system data")

        try:
            self.systemdata = []
            self._cli("show system")
            for line in self._read_long_output("show system"):
                self.systemdata.append(line)

            logging.debug("Grabbing system sw and hw versions")
            self._cli("show version")
            for line in self._read_long_output("show version"):
                self.systemdata.append(line)

        except PExpectError:
            # recurse on error
            self._switch_connect()
            return self._get_systemdata()

    ######################################
    # Internal port access helper methods
    ######################################
    # N.B. No parameter checking here, for speed reasons - if you're
    # calling this internal API then you should already have validated
    # things yourself! Equally, no post-set checks in here - do that
    # at the higher level.
    ######################################
    
    # Wrapper around connection.send - by default, expect() the same
    # text we've sent, to remove it from the output from the
    # switch. For the few cases where we don't need that, override
    # this using echo=False.
    # Horrible, but seems to work.
    def _cli(self, text, echo=True):
        self.connection.send(text + '\r')
        if echo:
            try:
                self.connection.expect(text)
            except (pexpect.EOF, pexpect.TIMEOUT):
                # Something went wrong; logout, log in and try again!
                logging.error("PEXPECT FAILURE, RECONNECT")
                self.errors.log_error_out(text)
                raise PExpectError("_cli failed on %s" % text)
            except:
                logging.error("Unexpected error: %s", sys.exc_info()[0])
                raise

if __name__ == "__main__":
#    p = CiscoSX300('10.172.2.52', 23)

    import optparse

    switch = 'vlandswitch02'
    parser = optparse.OptionParser()
    parser.add_option("--switch",
                      dest = "switch",
                      action = "store",
                      nargs = 1,
                      type = "string",
                      help = "specify switch to connect to for testing",
                      metavar = "<switch>")
    (opts, args) = parser.parse_args()
    if opts.switch:
        switch = opts.switch
        
    portname_base = "fa"
    # Text to match if we're on a SG-series switch, ports all called gi<number>
    sys_descr_re = re.compile('System Description.*Gigabit')

    def _port_name(number):
        return "%s%d" % (portname_base, number)
        
    logging.basicConfig(level = logging.DEBUG,
                        format = '%(asctime)s %(levelname)-8s %(message)s')
    p = CiscoSX300(switch, 23, debug = True)
    p.switch_connect('cisco', 'cisco', None)
    #buf = p._show_clock()
    #print "%s" % buf
    #buf = p._show_config()
    #p.dump_list(buf)

    print "System data:"
    p.dump_list(p.systemdata)
    for line in p.systemdata:
        match = sys_descr_re.match(line)
        if match:
            print 'Found an SG switch, using "gi" as port name prefix for testing'
            portname_base = "gi"

    if portname_base == "fa":
        print 'Found an SF switch, using "fa" as port name prefix for testing'

    print "Creating VLANs for testing:"
    for i in [ 2, 3, 4, 5, 20 ]:
        p.vlan_create(i)
        p.vlan_set_name(i, "test%d" % i)
        print "  %d (test%d)" % (i, i)

    #print "And dump config\n"
    #buf = p._show_config()
    #print "%s" % buf

    #print "Destroying VLAN 2\n"
    #p.vlan_destroy(2)

    #print "And dump config\n"
    #buf = p._show_config()
    #print "%s" % buf

    #print "Port names are:"
    #buf = p.switch_get_port_names()
    #p.dump_list(buf)

    #buf = p.vlan_get_name(25)
    #print "VLAN with tag 25 is called \"%s\"" % buf

    #p.vlan_set_name(35, "foo")
    #print "VLAN with tag 35 is called \"foo\""

    #buf = p.port_get_mode(_port_name(12))
    #print "Port %s is in %s mode" % (_port_name(12), buf)

    # Test access stuff
    print "Set %s to access mode" % _port_name(6)
    p.port_set_mode(_port_name(6), "access")
    print "Move %s to VLAN 2" % _port_name(6)
    p.port_set_access_vlan(_port_name(6), 2)
    buf = p.port_get_access_vlan(_port_name(6))
    print "Read from switch: %s is on VLAN %s" % (_port_name(6), buf)
    print "Move %s back to default VLAN 1" % _port_name(6)
    p.port_set_access_vlan(_port_name(6), 1)
    #print "And move %s back to a trunk port" % _port_name(6)
    #p.port_set_mode(_port_name(6), "trunk")
    #buf = p.port_get_mode(_port_name(6))
    #print "Port %s is in %s mode" % (_port_name(6), buf)

    # Test trunk stuff
    print "Set %s to trunk mode" % _port_name(2)
    p.port_set_mode(_port_name(2), "trunk")
    print "Add %s to VLAN 2" % _port_name(2)
    p.port_add_trunk_to_vlan(_port_name(2), 2)
    print "Add %s to VLAN 3" % _port_name(2)
    p.port_add_trunk_to_vlan(_port_name(2), 3)
    print "Add %s to VLAN 4" % _port_name(2)
    p.port_add_trunk_to_vlan(_port_name(2), 4)
    print "Read from switch: which VLANs is %s on?" % _port_name(2)
    buf = p.port_get_trunk_vlan_list(_port_name(2))
    p.dump_list(buf)

    print "Remove %s from VLANs 3,3,4" % _port_name(2)
    p.port_remove_trunk_from_vlan(_port_name(2), 3)
    p.port_remove_trunk_from_vlan(_port_name(2), 3)
    p.port_remove_trunk_from_vlan(_port_name(2), 4)
    print "Read from switch: which VLANs is %s on?" % _port_name(2)
    buf = p.port_get_trunk_vlan_list(_port_name(2))
    p.dump_list(buf)

 #   print "Adding lots of ports to VLANs"
 #   p.port_add_trunk_to_vlan(_port_name(1), 2)
 #   p.port_add_trunk_to_vlan(_port_name(3), 2)
 #   p.port_add_trunk_to_vlan(_port_name(5), 2)
 #   p.port_add_trunk_to_vlan(_port_name(7), 2)
 #   p.port_add_trunk_to_vlan(_port_name(9), 2)
 #   p.port_add_trunk_to_vlan(_port_name(11), 2)
 #   p.port_add_trunk_to_vlan(_port_name(13), 2)
 #   p.port_add_trunk_to_vlan(_port_name(15), 2)
 #   p.port_add_trunk_to_vlan(_port_name(17), 2)
 #   p.port_add_trunk_to_vlan(_port_name(19), 2)
 #   p.port_add_trunk_to_vlan(_port_name(21), 2)
 #   p.port_add_trunk_to_vlan(_port_name(23), 2)
 #   p.port_add_trunk_to_vlan(_port_name(4, 2)

    print "VLANs are:"
    buf = p.vlan_get_list()
    p.dump_list(buf)

#    print 'Restarting switch, to explicitly reset config'
#    p.switch_restart()

#    p.switch_save_running_config()
#    p._show_config()