util.py

#  Copyright 2014-2016 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.
#
#  Utility routines, including handling of API functions
#

import logging
import time
from errors import CriticalError, NotFoundError, InputError, ConfigError, SocketError

class VlanUtil:
    """VLANd utility functions"""

    def set_logging_level(self, level):
        loglevel = logging.CRITICAL
        if level == "ERROR":
            loglevel = logging.ERROR
        elif level == "WARNING":
            loglevel = logging.WARNING
        elif level == "INFO":
            loglevel = logging.INFO
        elif level == "DEBUG":
            loglevel = logging.DEBUG
        return loglevel

    def get_switch_driver(self, switch_name, config):
        logging.debug("Trying to find a driver for %s", switch_name)
        driver = config.switches[switch_name].driver
        logging.debug("Driver: %s", driver)
        module = __import__("drivers.%s" % driver, fromlist=[driver])
        class_ = getattr(module, driver)
        return class_(switch_name, debug = config.switches[switch_name].debug)

    def probe_switches(self, state):
        config = state.config
        ret = {}
        for switch_name in sorted(config.switches):
            logging.debug("Found switch %s:", switch_name)
            logging.debug("  Probing...")

            s = self.get_switch_driver(switch_name, config)
            s.switch_connect(config.switches[switch_name].username,
                             config.switches[switch_name].password,
                             config.switches[switch_name].enable_password)
            ret[switch_name] = 'Found %d ports: ' % len(s.switch_get_port_names())
            for name in s.switch_get_port_names():
                ret[switch_name] += '%s ' % name
            s.switch_disconnect()
            del s
        return ret

    # Simple helper wrapper for all the read-only database queries
    def perform_db_query(self, state, command, data):
        logging.debug('perform_db_query')
        logging.debug(command)
        logging.debug(data)
        ret = {}
        db = state.db
        try:
            if command == 'db.all_switches':
                ret = db.all_switches()
            elif command == 'db.all_ports':
                ret = db.all_ports()
            elif command == 'db.all_vlans':
                ret = db.all_vlans()
            elif command == 'db.all_trunks':
                ret = db.all_trunks()
            elif command == 'db.get_switch_by_id':
                ret = db.get_switch_by_id(data['switch_id'])
            elif command == 'db.get_switch_id_by_name':
                ret = db.get_switch_id_by_name(data['name'])
            elif command == 'db.get_switch_name_by_id':
                ret = db.get_switch_name_by_id(data['switch_id'])
            elif command == 'db.get_port_by_id':
                ret = db.get_port_by_id(data['port_id'])
            elif command == 'db.get_ports_by_switch':
                ret = db.get_ports_by_switch(data['switch_id'])
            elif command == 'db.get_port_by_switch_and_name':
                ret = db.get_port_by_switch_and_name(data['switch_id'], data['name'])
            elif command == 'db.get_port_by_switch_and_number':
                ret = db.get_port_by_switch_and_number(data['switch_id'], int(data['number']))
            elif command == 'db.get_current_vlan_id_by_port':
                ret = db.get_current_vlan_id_by_port(data['port_id'])
            elif command == 'db.get_base_vlan_id_by_port':
                ret = db.get_base_vlan_id_by_port(data['port_id'])
            elif command == 'db.get_ports_by_current_vlan':
                ret = db.get_ports_by_current_vlan(data['vlan_id'])
            elif command == 'db.get_ports_by_base_vlan':
                ret = db.get_ports_by_base_vlan(data['vlan_id'])
            elif command == 'db.get_ports_by_trunk':
                ret = db.get_ports_by_trunk(data['trunk_id'])
            elif command == 'db.get_vlan_by_id':
                ret = db.get_vlan_by_id(data['vlan_id'])
            elif command == 'db.get_vlan_tag_by_id':
                ret = db.get_vlan_tag_by_id(data['vlan_id'])
            elif command == 'db.get_vlan_id_by_name':
                ret = db.get_vlan_id_by_name(data['name'])
            elif command == 'db.get_vlan_id_by_tag':
                ret = db.get_vlan_id_by_tag(data['tag'])
            elif command == 'db.get_vlan_name_by_id':
                ret = db.get_vlan_name_by_id(data['vlan_id'])
            elif command == 'db.get_trunk_by_id':
                ret = db.get_trunk_by_id(data['trunk_id'])
            else:
                raise InputError("Unknown db_query command \"%s\"" % command)

        except (InputError, NotFoundError) as e:
            logging.error('perform_db_query(%s) got error %s', command, e)
            raise
        except ValueError as e:
            logging.error('perform_db_query(%s) got error %s', command, e)
            raise InputError("Invalid value in API call argument: %s" % e)
        except TypeError as e:
            logging.error('perform_db_query(%s) got error %s', command, e)
            raise InputError("Invalid type in API call argument: %s" % e)

        return ret

    # Simple helper wrapper for all the read-only daemon state queries
    def perform_daemon_query(self, state, command, data):
        logging.debug('perform_daemon_query')
        logging.debug(command)
        logging.debug(data)
        ret = {}
        try:
            if command == 'daemon.status':
                # data ignored
                ret['running'] = 'ok'
                ret['last_modified'] = state.db.get_last_modified_time()
            elif command == 'daemon.version':
                # data ignored
                ret['version'] = state.version
            elif command == 'daemon.statistics':
                ret['uptime'] = time.time() - state.starttime
            elif command == 'daemon.probe_switches':
                ret = self.probe_switches(state)
            elif command == 'daemon.shutdown':
                # data ignored
                ret['shutdown'] = 'Shutting down'
                state.running = False
            else:
                raise InputError("Unknown daemon_query command \"%s\"" % command)

        except (InputError, NotFoundError) as e:
            logging.error('perform_daemon_query(%s) got error %s', command, e)
            raise
        except ValueError as e:
            logging.error('perform_daemon_query(%s) got error %s', command, e)
            raise InputError("Invalid value in API call argument: %s" % e)
        except TypeError as e:
            logging.error('perform_daemon_query(%s) got error %s', command, e)
            raise InputError("Invalid type in API call argument: %s" % e)

        return ret

    # Helper wrapper for API functions modifying database state only
    def perform_db_update(self, state, command, data):
        logging.debug('perform_db_update')
        logging.debug(command)
        logging.debug(data)
        ret = {}
        db = state.db
        try:
            if command == 'db.create_switch':
                ret = db.create_switch(data['name'])
            elif command == 'db.create_port':
                try:
                    number = int(data['number'])
                except ValueError:
                    raise InputError("Invalid value for port number (%s) - must be numeric only!" % data['number'])
                ret = db.create_port(data['switch_id'], data['name'],
                                     number,
                                     state.default_vlan_id,
                                     state.default_vlan_id)
            elif command == 'db.create_trunk':
                ret = db.create_trunk(data['port_id1'], data['port_id2'])
            elif command == 'db.delete_switch':
                ret = db.delete_switch(data['switch_id'])
            elif command == 'db.delete_port':
                ret = db.delete_port(data['port_id'])
            elif command == 'db.set_port_is_locked':
                ret = db.set_port_is_locked(data['port_id'],
                                            data['is_locked'],
                                            data['lock_reason'])
            elif command == 'db.set_base_vlan':
                ret = db.set_base_vlan(data['port_id'], data['base_vlan_id'])
            elif command == 'db.delete_trunk':
                ret = db.delete_trunk(data['trunk_id'])
            else:
                raise InputError("Unknown db_update command \"%s\"" % command)

        except (InputError, NotFoundError) as e:
            logging.error('perform_db_update(%s) got error %s', command, e)
            raise
        except ValueError as e:
            logging.error('perform_db_update(%s) got error %s', command, e)
            raise InputError("Invalid value in API call argument: %s" % e)
        except TypeError as e:
            logging.error('perform_db_update(%s) got error %s', command, e)
            raise InputError("Invalid type in API call argument: %s" % e)

        return ret

    # Helper wrapper for API functions that modify both database state
    # and on-switch VLAN state
    def perform_vlan_update(self, state, command, data):
        logging.debug('perform_vlan_update')
        logging.debug(command)
        logging.debug(data)
        ret = {}

        try:
            # All of these are complex commands, so call helpers
            # rather than inline the code here
            if command == 'api.create_vlan':
                ret = self.create_vlan(state, data['name'], int(data['tag']), data['is_base_vlan'])
            elif command == 'api.delete_vlan':
                ret = self.delete_vlan(state, int(data['vlan_id']))
            elif command == 'api.set_port_mode':
                ret = self.set_port_mode(state, int(data['port_id']), data['mode'])
            elif command == 'api.set_current_vlan':
                ret = self.set_current_vlan(state, int(data['port_id']), int(data['vlan_id']))
            elif command == 'api.restore_base_vlan':
                ret = self.restore_base_vlan(state, int(data['port_id']))
            elif command == 'api.auto_import_switch':
                ret = self.auto_import_switch(state, data['switch'])
            else:
                raise InputError("Unknown query command \"%s\"" % command)

        except (InputError, NotFoundError) as e:
            logging.error('perform_vlan_update(%s) got error %s', command, e)
            raise
        except ValueError as e:
            logging.error('perform_vlan_update(%s) got error %s', command, e)
            raise InputError("Invalid value in API call argument: %s" % e)
        except TypeError as e:
            logging.error('perform_vlan_update(%s) got error %s', command, e)
            raise InputError("Invalid type in API call argument: %s" % e)

        return ret


    # Complex call
    # 1. create the VLAN in the DB
    # 2. Iterate through all switches:
    #    a. Create the VLAN
    #    b. Add the VLAN to all trunk ports (if needed)
    # 3. If all went OK, save config on all the switches
    #
    # The VLAN may already exist on some of the switches, that's
    # fine. If things fail, we attempt to roll back by rebooting
    # switches then removing the VLAN in the DB.
    def create_vlan(self, state, name, tag, is_base_vlan):

        logging.debug('create_vlan')
        db = state.db
        config = state.config

        # Check for tag == -1, i.e. use the next available tag
        if tag == -1:
            tag = db.find_lowest_unused_vlan_tag()
            logging.debug('create_vlan called with a tag of -1, found first unused tag %d', tag)

        # 1. Database record first
        try:
            logging.debug('Adding DB record first: name %s, tag %d, is_base_vlan %d', name, tag, is_base_vlan)
            vlan_id = db.create_vlan(name, tag, is_base_vlan)
            logging.debug('Added VLAN tag %d, name %s to the database, created VLAN ID %d', tag, name, vlan_id)
        except (InputError, NotFoundError):
            logging.debug('DB creation failed')
            raise

        # Keep track of which switches we've configured, for later use
        switches_done = []

        # 2. Now the switches
        try:
            for switch in db.all_switches():
                trunk_ports = []
                switch_name = switch.name
                try:
                    logging.debug('Adding new VLAN to switch %s', switch_name)
                    # Get the right driver
                    s = self.get_switch_driver(switch_name, config)
                    s.switch_connect(config.switches[switch_name].username,
                                     config.switches[switch_name].password,
                                     config.switches[switch_name].enable_password)

                    # Mark this switch as one we've touched, for
                    # either config saving or rollback below
                    switches_done.append(switch_name)

                    # 2a. Create the VLAN on the switch
                    s.vlan_create(tag)
                    s.vlan_set_name(tag, name)
                    logging.debug('Added VLAN tag %d, name %s to switch %s', tag, name, switch_name)

                    # 2b. Do we need to worry about trunk ports on this switch?
                    if 'TrunkWildCardVlans' in s.switch_get_capabilities():
                        logging.debug('This switch does not need special trunk port handling')
                    else:
                        logging.debug('This switch needs special trunk port handling')
                        trunk_ports = db.get_trunk_port_names_by_switch(switch.switch_id)
                        if trunk_ports is None:
                            logging.debug("But it has no trunk ports defined")
                            trunk_ports = []
                        else:
                            logging.debug('Found %d trunk_ports that need adjusting', len(trunk_ports))

                    # Modify any trunk ports as needed
                    for port in trunk_ports:
                        logging.debug('Adding VLAN tag %d, name %s to switch %s port %s', tag, name, switch_name, port)
                        s.port_add_trunk_to_vlan(port, tag)

                    # And now we're done with this switch
                    s.switch_disconnect()
                    del s

                except IOError as e:
                    logging.error('Failed to add VLAN %d to switch ID %d (%s): %s', tag, switch.switch_id, switch.name, e)
                    raise

        except IOError:
            # Bugger. Looks like one of the switch calls above
            # failed. To undo the changes safely, we'll need to reset
            # all the switches we managed to configure. This could
            # take some time!
            logging.error('create_vlan failed, resetting all switches to recover')
            for switch_name in switches_done:
                s = self.get_switch_driver(switch_name, config)
                s.switch_connect(config.switches[switch_name].username,
                                 config.switches[switch_name].password,
                                 config.switches[switch_name].enable_password)
                s.switch_restart() # Will implicitly also close the connection
                del s

            # Undo the database change
            logging.debug('Switch access failed. Deleting the new VLAN entry in the database')
            db.delete_vlan(vlan_id)
            raise

        # If we've got this far, things were successful. Save config
        # on all the switches so it will persist across reboots
        for switch_name in switches_done:
            s = self.get_switch_driver(switch_name, config)
            s.switch_connect(config.switches[switch_name].username,
                             config.switches[switch_name].password,
                             config.switches[switch_name].enable_password)
            s.switch_save_running_config()
            s.switch_disconnect()
            del s

        return (vlan_id, tag) # If we're successful

    # Complex call
    # 1. Check in the DB if there are any ports on the VLAN. Bail if so
    # 2. Iterate through all switches:
    #    a. Remove the VLAN from all trunk ports (if needed)
    #    b. Remove the VLAN
    # 3. If all went OK, save config on the switches
    # 4. Remove the VLAN in the DB
    #
    # If things fail, we attempt to roll back by rebooting switches.
    def delete_vlan(self, state, vlan_id):

        logging.debug('delete_vlan')
        db = state.db
        config = state.config

        # 1. Check for database records first
        logging.debug('Checking for ports using VLAN ID %d', vlan_id)
        vlan = db.get_vlan_by_id(vlan_id)
        if vlan is None:
            raise NotFoundError("VLAN ID %d does not exist" % vlan_id)
        vlan_tag = vlan.tag
        ports = db.get_ports_by_current_vlan(vlan_id)
        if ports is not None:
            raise InputError("Cannot delete VLAN ID %d when it still has %d ports" %
                             (vlan_id, len(ports)))
        ports = db.get_ports_by_base_vlan(vlan_id)
        if ports is not None:
            raise InputError("Cannot delete VLAN ID %d when it still has %d ports" %
                             (vlan_id, len(ports)))

        # Keep track of which switches we've configured, for later use
        switches_done = []

        # 2. Now the switches
        try:
            for switch in db.all_switches():
                switch_name = switch.name
                trunk_ports = []
                try:
                    # Get the right driver
                    s = self.get_switch_driver(switch_name, config)
                    s.switch_connect(config.switches[switch_name].username,
                                     config.switches[switch_name].password,
                                     config.switches[switch_name].enable_password)

                    # Mark this switch as one we've touched, for
                    # either config saving or rollback below
                    switches_done.append(switch_name)

                    # 2a. Do we need to worry about trunk ports on this switch?
                    if 'TrunkWildCardVlans' in s.switch_get_capabilities():
                        logging.debug('This switch does not need special trunk port handling')
                    else:
                        logging.debug('This switch needs special trunk port handling')
                        trunk_ports = db.get_trunk_port_names_by_switch(switch.switch_id)
                        if trunk_ports is None:
                            logging.debug("But it has no trunk ports defined")
                            trunk_ports = []
                        else:
                            logging.debug('Found %d trunk_ports that need adjusting', len(trunk_ports))

                    # Modify any trunk ports as needed
                    for port in trunk_ports:
                        s.port_remove_trunk_from_vlan(port, vlan_tag)
                        logging.debug('Removed VLAN tag %d from switch %s port %s', vlan_tag, switch_name, port)

                    # 2b. Remove the VLAN from the switch
                    logging.debug('Removing VLAN tag %d from switch %s', vlan_tag, switch_name)
                    s.vlan_destroy(vlan_tag)
                    logging.debug('Removed VLAN tag %d from switch %s', vlan_tag, switch_name)

                    # And now we're done with this switch
                    s.switch_disconnect()
                    del s

                except IOError:
                    raise

        except IOError:
            # Bugger. Looks like one of the switch calls above
            # failed. To undo the changes safely, we'll need to reset
            # all the switches we managed to configure. This could
            # take some time!
            logging.error('delete_vlan failed, resetting all switches to recover')
            for switch_name in switches_done:
                s = self.get_switch_driver(switch_name, config)
                s.switch_connect(config.switches[switch_name].username,
                                 config.switches[switch_name].password,
                                 config.switches[switch_name].enable_password)
                s.switch_restart() # Will implicitly also close the connection
                del s
                raise

        # 3. If we've got this far, things were successful. Save
        # config on all the switches so it will persist across reboots
        for switch_name in switches_done:
            s = self.get_switch_driver(switch_name, config)
            s.switch_connect(config.switches[switch_name].username,
                             config.switches[switch_name].password,
                             config.switches[switch_name].enable_password)
            s.switch_save_running_config()
            s.switch_disconnect()
            del s

        # 4. Finally, remove the VLAN in the DB
        try:
            logging.debug('Removing DB record: VLAN ID %d', vlan_id)
            vlan_id = db.delete_vlan(vlan_id)
            logging.debug('Removed VLAN ID %d from the database OK', vlan_id)
        except (InputError, NotFoundError):
            logging.debug('DB deletion failed')
            raise

        return vlan_id # If we're successful

    # Complex call, depends on existing state a lot
    # 1. Check validity of inputs
    # 2. Switch mode and other config on the port.
    #    a. If switching trunk->access, remove all trunk VLANs from it
    #       (if needed) and switch back to the base VLAN for the
    #       port. Next, switch to access mode.
    #    b. If switching access->trunk, switch back to the base VLAN
    #       for the port. Next, switch mode. Then add all trunk VLANs
    #       to it (if needed)
    # 3. If all went OK, save config on the switch
    # 4. Change details of the port in the DB
    #
    # If things fail, we attempt to roll back by rebooting the switch
    def set_port_mode(self, state, port_id, mode):

        logging.debug('set_port_mode')
        db = state.db
        config = state.config

        # 1. Sanity-check inputs
        if mode != 'access' and mode != 'trunk':
            raise InputError("Port mode '%s' is not a valid option: try 'access' or 'trunk'" % mode)
        port = db.get_port_by_id(port_id)
        if port is None:
            raise NotFoundError("Port ID %d does not exist" % port_id)
        if port.is_locked:
            raise InputError("Port ID %d is locked" % port_id)
        if mode == 'trunk' and port.is_trunk:
            raise InputError("Port ID %d is already in trunk mode" % port_id)
        if mode == 'access' and not port.is_trunk:
            raise InputError("Port ID %d is already in access mode" % port_id)
        base_vlan_tag = db.get_vlan_tag_by_id(port.base_vlan_id)

        # Get the right driver
        switch_name = db.get_switch_name_by_id(port.switch_id)
        s = self.get_switch_driver(switch_name, config)

        # 2. Now start configuring the switch
        try:
            s.switch_connect(config.switches[switch_name].username,
                             config.switches[switch_name].password,
                             config.switches[switch_name].enable_password)
        except:
            logging.debug('Failed to talk to switch %s!', switch_name)
            raise

        try:
            if port.is_trunk:
                # 2a. We're going from a trunk port to an access port
                if 'TrunkWildCardVlans' in s.switch_get_capabilities():
                    logging.debug('This switch does not need special trunk port handling')
                else:
                    logging.debug('This switch needs special trunk port handling')
                    vlans = s.port_get_trunk_vlan_list(port.name)
                    if vlans is None:
                        logging.debug("But it has no VLANs defined on port %s", port.name)
                        vlans = []
                    else:
                        logging.debug('Found %d vlans that may need dropping on port %s', len(vlans), port.name)

                    for vlan in vlans:
                        if vlan != state.config.vland.default_vlan_tag:
                            s.port_remove_trunk_from_vlan(port.name, vlan)

                s.port_set_mode(port.name, "access")

            else:
                # 2b. We're going from an access port to a trunk port
                s.port_set_access_vlan(port.name, base_vlan_tag)
                s.port_set_mode(port.name, "trunk")
                if 'TrunkWildCardVlans' in s.switch_get_capabilities():
                    logging.debug('This switch does not need special trunk port handling')
                else:
                    vlans = db.all_vlans()
                    for vlan in vlans:
                        if vlan.tag != state.config.vland.default_vlan_tag:
                            s.port_add_trunk_to_vlan(port.name, vlan.tag)

        except IOError:
            logging.error('set_port_mode failed, resetting switch to recover')
            # Bugger. Looks like one of the switch calls above
            # failed. To undo the changes safely, we'll need to reset
            # all the config on this switch
            s.switch_restart() # Will implicitly also close the connection
            del s
            raise

        # 3. All seems to have worked so far!
        s.switch_save_running_config()
        s.switch_disconnect()
        del s

        # 4. And update the DB
        db.set_port_mode(port_id, mode)

        return port_id # If we're successful

    # Complex call, updating both DB and switch state
    # 1. Check validity of inputs
    # 2. Update the port config on the switch
    # 3. If all went OK, save config on the switch
    # 4. Change details of the port in the DB
    #
    # If things fail, we attempt to roll back by rebooting the switch
    def set_current_vlan(self, state, port_id, vlan_id):

        logging.debug('set_current_vlan')
        db = state.db
        config = state.config

        # 1. Sanity checks!
        port = db.get_port_by_id(port_id)
        if port is None:
            raise NotFoundError("Port ID %d does not exist" % port_id)
        if port.is_locked:
            raise InputError("Port ID %d is locked" % port_id)
        if port.is_trunk:
            raise InputError("Port ID %d is not an access port" % port_id)

        vlan = db.get_vlan_by_id(vlan_id)
        if vlan is None:
            raise NotFoundError("VLAN ID %d does not exist" % vlan_id)

        # Get the right driver
        switch_name = db.get_switch_name_by_id(port.switch_id)
        s = self.get_switch_driver(switch_name, config)

        # 2. Now start configuring the switch
        try:
            s.switch_connect(config.switches[switch_name].username,
                             config.switches[switch_name].password,
                             config.switches[switch_name].enable_password)
        except:
            logging.debug('Failed to talk to switch %s!', switch_name)
            raise

        try:
            s.port_set_access_vlan(port.name, vlan.tag)
        except IOError:
            logging.error('set_current_vlan failed, resetting switch to recover')
            # Bugger. Looks like one of the switch calls above
            # failed. To undo the changes safely, we'll need to reset
            # all the config on this switch
            s.switch_restart() # Will implicitly also close the connection
            del s
            raise

        # 3. All seems to have worked so far!
        s.switch_save_running_config()
        s.switch_disconnect()
        del s

        # 4. And update the DB
        db.set_current_vlan(port_id, vlan_id)

        return port_id # If we're successful

    # Complex call, updating both DB and switch state
    # 1. Check validity of input
    # 2. Update the port config on the switch
    # 3. If all went OK, save config on the switch
    # 4. Change details of the port in the DB
    #
    # If things fail, we attempt to roll back by rebooting the switch
    def restore_base_vlan(self, state, port_id):

        logging.debug('restore_base_vlan')
        db = state.db
        config = state.config

        # 1. Sanity checks!
        port = db.get_port_by_id(port_id)
        if port is None:
            raise NotFoundError("Port ID %d does not exist" % port_id)
        if port.is_trunk:
            raise InputError("Port ID %d is not an access port" % port_id)
        if port.is_locked:
            raise InputError("Port ID %d is locked" % port_id)

        # Bail out early if we're *already* on the base VLAN. This is
        # not an error
        if port.current_vlan_id == port.base_vlan_id:
            return port_id

        vlan = db.get_vlan_by_id(port.base_vlan_id)

        # Get the right driver
        switch_name = db.get_switch_name_by_id(port.switch_id)
        s = self.get_switch_driver(switch_name, config)

        # 2. Now start configuring the switch
        try:
            s.switch_connect(config.switches[switch_name].username,
                             config.switches[switch_name].password,
                             config.switches[switch_name].enable_password)
        except:
            logging.debug('Failed to talk to switch %s!', switch_name)
            raise

        try:
            s.port_set_access_vlan(port.name, vlan.tag)
        except IOError:
            logging.error('restore_base_vlan failed, resetting switch to recover')
            # Bugger. Looks like one of the switch calls above
            # failed. To undo the changes safely, we'll need to reset
            # all the config on this switch
            s.switch_restart() # Will implicitly also close the connection
            del s
            raise

        # 3. All seems to have worked so far!
        s.switch_save_running_config()
        s.switch_disconnect()
        del s

        # 4. And update the DB
        db.set_current_vlan(port_id, port.base_vlan_id)

        return port_id # If we're successful

    # Complex call, updating both DB and switch state
    # * Check validity of input
    # * Read all the config from the switch (switch, ports, VLANs)
    # * Create initial DB entries to match each of those
    # * Merge VLANs across all switches
    # * Set up ports appropriately
    #
    def auto_import_switch(self, state, switch_name):

        logging.debug('auto_import_switch')
        db = state.db
        config = state.config

        port_vlans = {}

        # 1. Sanity checks!
        switch_id = db.get_switch_id_by_name(switch_name)
        if switch_id is not None:
            raise InputError("Switch name %s already exists in the DB (ID %d)" % (switch_name, switch_id))

        if not switch_name in config.switches:
            raise NotFoundError("Switch name %s not defined in config" % switch_name)

        # 2. Now start reading config from the switch
        try:
            s = self.get_switch_driver(switch_name, config)
            s.switch_connect(config.switches[switch_name].username,
                             config.switches[switch_name].password,
                             config.switches[switch_name].enable_password)
        except:
            logging.debug('Failed to talk to switch %s!', switch_name)
            raise

        # DON'T create the switch record in the DB first - we'll want
        # to create VLANs on *other* switches, and it's easier to do
        # that before we've added our new switch

        new_vlan_tags = []

        # Grab the VLANs defined on this switch
        vlan_tags = s.vlan_get_list()

        logging.debug('  found %d vlans on the switch', len(vlan_tags))

        for vlan_tag in vlan_tags:
            vlan_name = s.vlan_get_name(vlan_tag)

            # If a VLAN is already in the database, then that's easy -
            # we can just ignore it. However, we have to check that
            # there is not a different name for the existing VLAN tag
            # - bail out if so... UNLESS we're looking at the default
            # VLAN
            #
            # If this VLAN tag is not already in the DB, we'll need to
            # add it there and to all the other switches (and their
            # trunk ports!) too.
            vlan_id = db.get_vlan_id_by_tag(vlan_tag)
            if vlan_id != state.default_vlan_id:
                if vlan_id is not None:
                    vlan_db_name = db.get_vlan_name_by_id(vlan_id)
                    if vlan_name != vlan_db_name:
                        raise InputError("Can't add VLAN tag %d (name %s) for this switch - VLAN tag %d already exists in the database, but with a different name (%s)" % (vlan_tag, vlan_name, vlan_tag, vlan_db_name))

                else:
                    # OK, we'll need to set up the new VLAN now. It can't
                    # be a base VLAN - switches don't have such a concept!
                    # Rather than create individually here, add to a
                    # list. *Only* once we've worked through all the
                    # switch's VLANs successfully (checking for existing
                    # records and possible clashes!) should we start
                    # committing changes
                    new_vlan_tags.append(vlan_tag)

        # Now create the VLAN DB entries
        for vlan_tag in new_vlan_tags:
            vlan_name = s.vlan_get_name(vlan_tag)
            vlan_id = self.create_vlan(state, vlan_name, vlan_tag, False)

        # *Now* add this switch itself to the database, after we've
        # worked on all the other switches
        switch_id = db.create_switch(switch_name)

        # And now the ports
        trunk_ports = []
        ports = s.switch_get_port_names()
        logging.debug('  found %d ports on the switch', len(ports))
        for port_name in ports:
            logging.debug('  trying to import port %s', port_name)
            port_id = None
            port_mode = s.port_get_mode(port_name)
            port_number = s.port_map_name_to_number(port_name)
            if port_mode == 'access':
                # Access ports are easy - just create the port, and
                # set both the current and base VLANs to the current
                # VLAN on the switch. We'll end up changing this after
                # import if needed.
                port_vlans[port_name] = (s.port_get_access_vlan(port_name),)
                port_vlan_id = db.get_vlan_id_by_tag(port_vlans[port_name][0])
                port_id = db.create_port(switch_id, port_name, port_number,
                                         port_vlan_id, port_vlan_id)
                logging.debug('    access port, VLAN %d', int(port_vlans[port_name][0]))
                # Nothing further needed
            elif port_mode == 'trunk':
                logging.debug('    trunk port, VLANs:')
                # Trunk ports are a little more involved. First,
                # create the port in the DB, setting the VLANs to the
                # first VLAN found on the trunk port. This will *also*
                # be in access mode by default, and unlocked.
                port_vlans[port_name] = s.port_get_trunk_vlan_list(port_name)
                logging.debug(port_vlans[port_name])
                if port_vlans[port_name] == [] or port_vlans[port_name] is None or 'ALL' in port_vlans[port_name]:
                    port_vlans[port_name] = (state.config.vland.default_vlan_tag,)
                port_vlan_id = db.get_vlan_id_by_tag(port_vlans[port_name][0])
                port_id = db.create_port(switch_id, port_name, port_number,
                                         port_vlan_id, port_vlan_id)
                # Append to a list of trunk ports that we will need to
                # modify once we're done
                trunk_ports.append(port_id)
            else:
                # We've found a port mode we don't want, e.g. the
                # "dynamic auto" on a Cisco Catalyst. Handle that here
                # - tell the switch to set that port to access and
                # handle accordingly.
                s.port_set_mode(port_name, 'access')
                port_vlans[port_name] = (s.port_get_access_vlan(port_name),)
                port_vlan_id = db.get_vlan_id_by_tag(port_vlans[port_name][0])
                port_id = db.create_port(switch_id, port_name, port_number,
                                         port_vlan_id, port_vlan_id)
                logging.debug('    Found port in %s mode', port_mode)
                logging.debug('    Forcing to access mode, VLAN %d', int(port_vlans[port_name][0]))
                port_mode = "access"

            logging.debug("  Added port %s, got port ID %d", port_name, port_id)

            db.set_port_mode(port_id, port_mode)

        # Make sure this switch has all the VLANs we need
        for vlan in db.all_vlans():
            if vlan.tag != state.config.vland.default_vlan_tag:
                if not vlan.tag in vlan_tags:
                    logging.debug("Adding VLAN tag %d to this switch", vlan.tag)
                    s.vlan_create(vlan.tag)
                    s.vlan_set_name(vlan.tag, vlan.name)

        # Now, on each trunk port on the switch, we need to add all
        # the VLANs already configured across our system
        if not 'TrunkWildCardVlans' in s.switch_get_capabilities():
            for port_id in trunk_ports:
                port = db.get_port_by_id(port_id)

                for vlan in db.all_vlans():
                    if vlan.vlan_id != state.default_vlan_id:
                        if not vlan.tag in port_vlans[port.name]:
                            logging.debug("Adding allowed VLAN tag %d to trunk port %s", vlan.tag, port.name)
                            s.port_add_trunk_to_vlan(port.name, vlan.tag)

        # Done with this switch \o/
        s.switch_save_running_config()
        s.switch_disconnect()
        del s

        ret = {}
        ret['switch_id'] = switch_id
        ret['num_ports_added'] = len(ports)
        ret['num_vlans_added'] = len(new_vlan_tags)
        return ret # If we're successful