Vland/drivers/Mellanox.py

#! /usr/bin/python

#  Copyright 2018 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 sys
import re
import pexpect

# Mellanox MLNX-OS driver
# Developed and tested against the SN2100 in the Linaro LAVA lab

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 Mellanox(SwitchDriver):

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

    _capabilities = [
        'TrunkWildCardVlans' # Trunk ports are on all VLANs by
                             # default, so we shouldn't need to
                             # bugger with them
    ]

    # Regexp of expected hardware information - fail if we don't see
    # this
    _expected_descr_re = re.compile(r'MLNX-OS')

    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
    ################################

    # Save the current running config into flash - we want config to
    # remain across reboots
    def switch_save_running_config(self):
        try:
            self._configure()
            self._cli("configuration write")
            self._end_configure()
        except (PExpectError, pexpect.EOF):
            # 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 noconfirm")

    # List the capabilities of the switch (and driver) - some things
    # make no sense to abstract. Returns a dict of strings, each one
    # describing an extra feature that that higher levels may care
    # about
    def switch_get_capabilities(self):
        return self._capabilities

    ################################
    ### 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 %d" % tag)
            self._end_vlan()
            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 name %s" % (tag, name))
            self._end_configure()

            # This switch *might* have problems if we drive it too quickly? At
            # least one instance of set_name()/get_name() not working. This
            # might help?
            self._delay()

            # And retry around here
            retries = 5
            read_name = None
            while (retries > 0 and read_name is None):
                # Validate it happened
                read_name = self.vlan_get_name(tag)
                retries -= 1

            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(r'^ *(\d+)')

            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):

        try:
            logging.debug("Grabbing the name of VLAN %d", tag)
            name = None

            # Ugh, the output here is messy. VLAN names can include spaces, and
            # there are no delimiters in the output, e.g.:
            # VLAN    Name                    Ports
            # ----    -----------             --------------------------------------
            # 1       default                 Eth1/1/1, Eth1/1/2, Eth1/2, Eth1/3/1, Eth1/3/2,
            #                                 Eth1/4, Eth1/5, Eth1/6, Eth1/7, Eth1/8,
            #                                 Eth1/10, Eth1/12, Eth1/13, Eth1/14, Eth1/15,
            #                                 Eth1/16
            # 102     mdev testing
            # 103     vpp 1 performance testing       Eth1/1/3, Eth1/9
            # 104     vpp 2 performance testing       Eth1/1/4, Eth1/11
            #
            # Simplest strategy:
            # 1. Match on a leading number and grab all the text after it
            # 2. Drop anything starting with "Eth" to EOL
            # 3. Strip leading and trailing whitespace
            #
            # Not perfect, but it'll have to do. Anybody including "Eth" in a
            # VLAN name deserves to lose...

            regex = re.compile(r'^ *\d+\s+(.+)')
            self._cli("show vlan id %d" % tag)
            for line in self._read_long_output("show vlan id"):
                match = regex.match(line)
                if match:
                    name = re.sub(r'Eth.*$',"",match.group(1)).strip()
            if name is None:
                logging.debug("vlan_get_name: did not find a name")
            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 ethernet %s" % port)
            self._cli("switchport mode %s" % mode)
            if mode == "trunk":
                # Put the new trunk port on all VLANs
                self._cli("switchport trunk allowed-vlan all")
            self._end_interface()
            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)

            # And cache the result
            self._port_modes[port] = mode

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

    # 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 ethernet %s" % port)
            self._cli("switchport access vlan %d" % tag)
            self._end_interface()
            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 ethernet %s" % port)
            self._cli("switchport trunk allowed-vlan add %d" % tag)
            self._end_interface()
            self._end_configure()

            # Validate it happened
            read_vlans = self.port_get_trunk_vlan_list(port)
            for vlan in read_vlans:
                if vlan == tag or vlan == "ALL":
                    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 ethernet %s" % port)
            self._cli("switchport trunk allowed-vlan remove %d" % tag)
            self._end_interface()
            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)
        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(r'^Eth%s\s+access\s+(\d+)' % port)

        try:
            self._cli("show interfaces switchport")
            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_start = re.compile(r'^Eth%s\s+trunk\s+N/A\s+(.*)' % port)
        regex_continue = re.compile(r'^(\d.*)')

        try:
            self._cli("show interfaces switchport")

            # Complex parsing work - VLAN list may extend over several lines, e.g.:
            #
            # Eth1/16         trunk        N/A                1, 102, 103, 104, 1000, 1001, 1002
            #                                                 1003, 1004
            #
            in_match = False
            vlan_text = ''

            for line in self._read_long_output("show interfaces switchport"):
                if in_match:
                    match = regex_continue.match(line)
                    if match:
                        vlan_text += ', ' # Make a consistently-formed list
                        vlan_text += match.group(1)
                    else:
                        in_match = False
                if not in_match:
                    match = regex_start.match(line)
                    if match:
                        vlan_text += match.group(1)
                        in_match = True

            vlans = self._parse_vlan_list(vlan_text)
            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.logger)
        self._login()

        # Avoid paged output as much as possible
        self._cli("terminal length 999")
        # Don't do silly things with ANSI codes
        self._cli("terminal type dumb")
        # and disable auto-logout after delay
        self._cli("no cli session auto-logout")

        # 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(r'Product name:\s*(\S+)')
        sn_regex = re.compile(r'System serial num:\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)

        logging.debug("serial number is %s", self.serial_number)
        logging.debug("system description is %s", descr)

        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)
        if self._username is not None:
            self.connection.expect("login:")
            self._cli("%s" % self._username)
        if self._password is not None:
            self.connection.expect("Password:")
            self._cli("%s" % self._password, False)
        while True:
            index = self.connection.expect(['User:', 'Password:', 'Login incorrect', 'authentication failed', r'(.*?)(#|>)'])
            if index != 4: # Any other means: failed to log in!
                logging.error("Login failure: index %d\n", index)
                logging.error("Login failure: %s\n", self.connection.match.before)
                raise IOError

            # else

            # Add a couple of newlines to get past the "last login" etc. junk
            self._cli("")
            self._cli("")
            self.connection.expect(r'^(.*?) (#|>)')
            self._prompt_name = re.escape(self.connection.match.group(1).strip())
            logging.info("Got outer prompt \"%s\"", self._prompt_name)
            if self.connection.match.group(2) == ">":
                # Need to enter "enable" mode too
                self._cli("enable")
                if self._enable_password is not None:
                    self.connection.expect("Password:")
                    self._cli("%s" % self._enable_password, False)
                    index = self.connection.expect(['Password:', 'Login incorrect', 'authentication failed', r'(.*) *(#|>)'])
                    if index != 3: # Any other means: failed to log in!
                        logging.error("Enable password failure: %s\n", self.connection.match)
                        raise IOError
                    self._cli("")
                    self._cli("")
                    self.connection.expect(r'^(.*?) (#|>)')
                    self._prompt_name = re.escape(self.connection.match.group(1).strip())
                    logging.info("Got enable prompt \"%s\"", self._prompt_name)
                return 0

    def _logout(self):
        logging.debug("Logging out")
        self._cli("quit", False)
        try:
            self.connection.expect("Would you like to save them now")
            self._cli("n")
        except (pexpect.EOF):
            pass
        self.connection.close(True)

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

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

    def _end_interface(self):
        self._cli("exit")

    def _end_vlan(self):
        self._cli("exit")

    def _read_long_output(self, text):
        longbuf = []
        prompt = self._prompt_name + r'\s*#'
        while True:
            try:
                index = self.connection.expect([r'lines \d+-\d+', prompt])
                if index == 0: # "lines 45-50"
                    for line in self.connection.before.split('\r\n'):
                        longbuf.append(line.strip())
                    self._cli(' ', False)
                elif index == 1: # Back to a prompt, says output is finished
                    break
            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

        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 = []

        # Look for "Eth1" at the beginning of the output lines to just
        # match lines with interfaces - they have names like
        # "Eth1/15". We do not care about Link Aggregation Groups (lag)
        # here.
        regex = re.compile(r'^Eth(\S+)')

        try:
            self._cli("show interfaces switchport")
            for line in self._read_long_output("show interfaces switchport"):
                match = regex.match(line)
                if match:
                    interface = match.group(1)
                    interfaces.append(interface)
                    self._port_numbers[interface] = len(interfaces)
            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()

    # 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('Switchport mode: (.*)')

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

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

    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 sw and hw versions")

        try:
            self._systemdata = []
            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()

    # Borrowed from the Catalyst driver. Over-complex for our needs here, but
    # it's already tested and will do the job.
    def _parse_vlan_list(self, inputdata):
        vlans = []

        if inputdata == "ALL":
            return ["ALL"]
        elif inputdata == "NONE":
            return []
        elif inputdata == "":
            return []
        else:
            # Parse the complex list
            groups = inputdata.split(',')
            for group in groups:
                subgroups = group.split('-')
                if len(subgroups) == 1:
                    vlans.append(int(subgroups[0]))
                elif len(subgroups) == 2:
                    for i in range (int(subgroups[0]), int(subgroups[1]) + 1):
                        vlans.append(i)
                else:
                    logging.debug("Can't parse group \"" + group + "\"")

        return vlans

    # 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__":

    # Simple test harness - exercise the main working functions above to verify
    # they work. This does *NOT* test really disruptive things like "save
    # running-config" and "reload" - test those by hand.

    import optparse

    switch = '172.27.16.6'
    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

    logging.basicConfig(level = logging.DEBUG,
                        format = '%(asctime)s %(levelname)-8s %(message)s')
    p = MlnxOS(switch, 23, debug=False)
    p.switch_connect('admin', 'admin', None)

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

    buf = p.vlan_get_name(102)
    print "VLAN 102 is named \"%s\"" % buf

    print "Create VLAN 1003"
    p.vlan_create(1003)

    buf = p.vlan_get_name(1003)
    print "VLAN 1003 is named \"%s\"" % buf

    print "Set name of VLAN 1003 to test333"
    p.vlan_set_name(1003, "test333")

    buf = p.vlan_get_name(1003)
    print "VLAN 1003 is named \"%s\"" % buf

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

    print "Destroy VLAN 1003"
    p.vlan_destroy(1003)

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

    buf = p.port_get_mode("1/15")
    print "Port 1/15 is in %s mode" % buf

    buf = p.port_get_mode("1/16")
    print "Port 1/16 is in %s mode" % buf

    # Test access stuff
    print "Set 1/15 to access mode"
    p.port_set_mode("1/15", "access")

    print "Move 1/15 to VLAN 4"
    p.port_set_access_vlan("1/15", 4)

    buf = p.port_get_access_vlan("1/15")
    print "Read from switch: 1/15 is on VLAN %s" % buf

    print "Move 1/15 back to VLAN 1"
    p.port_set_access_vlan("1/15", 1)

    print "Create VLAN 1002"
    p.vlan_create(1002)

    print "Create VLAN 1003"
    p.vlan_create(1003)

    print "Create VLAN 1004"
    p.vlan_create(1004)

    # Test access stuff
    print "Set 1/15 to trunk mode"
    p.port_set_mode("1/15", "trunk")
    print "Read from switch: which VLANs is 1/15 on?"
    buf = p.port_get_trunk_vlan_list("1/15")
    p.dump_list(buf)

    # The adds below are NOOPs in effect on this switch - no filtering
    # for "trunk" ports
    print "Add 1/15 to VLAN 1002"
    p.port_add_trunk_to_vlan("1/15", 1002)
    print "Add 1/15 to VLAN 1003"
    p.port_add_trunk_to_vlan("1/15", 1003)
    print "Add 1/15 to VLAN 1004"
    p.port_add_trunk_to_vlan("1/15", 1004)
    print "Read from switch: which VLANs is 1/15 on?"
    buf = p.port_get_trunk_vlan_list("1/15")
    p.dump_list(buf)

    # And the same for removals here
    p.port_remove_trunk_from_vlan("1/15", 1003)
    p.port_remove_trunk_from_vlan("1/15", 1002)
    p.port_remove_trunk_from_vlan("1/15", 1004)
    print "Read from switch: which VLANs is 1/15 on?"
    buf = p.port_get_trunk_vlan_list("1/15")
    p.dump_list(buf)

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

#    p.switch_save_running_config()

#    p.switch_disconnect()
#    p._show_config()