util.py

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

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

    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)

    def probe_switches(self, state):
        config = state.config
        ret = {}
        for switch_name in sorted(config.switches):
            print "Found switch %s:" % (switch_name)
            print "  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 switch %s with %d ports' % (switch_name, len(s.switch_get_port_names()))
            s.switch_disconnect()
            del(s)

    # Simple helper wrapper for all the read-only database queries
    def perform_db_query(self, state, command, data):
        print 'perform_db_query'
        print command
        print 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.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_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_vlan_by_id':
                ret = db.get_vlan_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'])
            else:
                raise InputError("Unknown db_query command \"%s\"" % command)

        except InputError:
            raise

        except:
            raise InputError("Invalid input in query")

        return ret

    # Simple helper wrapper for all the read-only daemon state queries
    def perform_daemon_query(self, state, command, data):
        print 'perform_daemon_query'
        print command
        print data
        ret = {}
        try:
            if command == 'daemon.status':
                # data ignored
                ret['running'] = 'ok'
            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)
            else:
                raise InputError("Unknown daemon_query command \"%s\"" % command)

        except InputError:
            raise

        except:
            raise InputError("Invalid input in query")

        return ret

    # Helper wrapper for API functions modifying database state only
    def perform_db_update(self, state, command, data):
        print 'perform_db_update'
        print command
        print data
        ret = {}
        db = state.db
        try:
            if command == 'db.create_switch':
                ret = db.create_switch(data['name'])
            elif command == 'db.create_port':
                ret = db.create_port(data['switch_id'], data['name'],
                                     state.default_vlan_id,
                                     state.default_vlan_id)
            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'])
            elif command == 'db.set_base_vlan':
                ret = db.set_base_vlan(data['port_id'], data['base_vlan_id'])
            else:
                raise InputError("Unknown db_update command \"%s\"" % command)

        except InputError:
            raise

        except:
            raise InputError("Invalid input in query")

        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):
        print 'perform_vlan_update'
        print command
        print 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 as e:
            print 'got error %s' % e
            raise

        except:
            raise InputError("Invalid input in query")

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

        print 'create_vlan'
        db = state.db
        config = state.config

        # 1. Database record first
        try:
            print '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)
            print 'Added VLAN tag %d, name %s to the database, created VLAN ID %d' % (tag, name, vlan_id)
        except InputError:
            print '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:
                    print '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)
                    print '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():
                        print 'This switch does not need special trunk port handling'
                    else:
                        print 'This switch needs special trunk port handling'
                        trunk_ports = db.get_trunk_port_names_by_switch(switch.switch_id)
                        if trunk_ports is None:
                            print "But it has no trunk ports defined"
                            trunk_ports = []
                        else:
                            print 'Found %d trunk_ports that need adjusting' % len(trunk_ports)

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

                    # 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!
            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
            print '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 # 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):

        print 'delete_vlan'
        db = state.db
        config = state.config

        # 1. Check for database records first
        print 'Checking for ports using VLAN ID %d' % vlan_id
        vlan = db.get_vlan_by_id(vlan_id)
        if vlan is None:
            raise InputError("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():
                        print 'This switch does not need special trunk port handling'
                    else:
                        print 'This switch needs special trunk port handling'
                        trunk_ports = db.get_trunk_port_names_by_switch(switch.switch_id)
                        if trunk_ports is None:
                            print "But it has no trunk ports defined"
                            trunk_ports = []
                        else:
                            print '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)
                        print 'Removed VLAN tag %d from switch %s port %s' % (vlan_tag, switch_name, port)

                    # 2b. Remove the VLAN from the switch
                    print 'Removing VLAN tag %d from switch %s' % (vlan_tag, switch_name)
                    s.vlan_destroy(vlan_tag)
                    print '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!
            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:
            print 'Removing DB record: VLAN ID %d' % vlan_id
            vlan_id = db.delete_vlan(vlan_id)
            print 'Removed VLAN ID %d from the database OK' % vlan_id
        except InputError:
            print '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):

        print 'set_port_mode'
        db = state.db
        config = state.config

        # 1. Sanity-check inputs
        if mode is not 'access' and mode is not '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 InputError("Port ID %d does not exist" % port_id)
        if mode == 'trunk' and port.is_trunk:
            raise InputError("Port ID %d is already in trunk mode")
        if mode == 'access' and not port.is_trunk:
            raise InputError("Port ID %d is already in access mode")
        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:
            print '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():
                    print 'This switch does not need special trunk port handling'
                else:
                    print 'This switch needs special trunk port handling'
                    vlans = s.port_get_trunk_vlan_list(port.name)
                    if vlans is None:
                        print "But it has no VLANs defined on port %s" % port.name
                        vlans = []
                    else:
                        print 'Found %d vlans that need dropping on port %s' % (len(vlans), port.name)

                    for vlan in vlans:
                        s.port_remove_trunk_from_vlan(port.name, vlan)
                    s.port_add_trunk_to_vlan(port.name, base_vlan_tag)
                    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():
                    print 'This switch does not need special trunk port handling'
                else:
                    vlans = db.all_vlans()
                    for vlan in vlans:
                        s.port_add_trunk_to_vlan(port.name, vlan.tag)

        except IOError:
            # 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):

        print '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 InputError("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)

        vlan = db.get_vlan_by_id(vlan_id)
        if vlan is None:
            raise InputError("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:
            print 'Failed to talk to switch %s!' % switch_name
            raise

        try:
            s.port_set_access_vlan(port.name, vlan.tag)
        except IOError:
            # 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):

        print '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 InputError("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:
            print 'Failed to talk to switch %s!' % switch_name
            raise

        try:
            s.port_set_access_vlan(port.name, vlan.tag)
        except IOError:
            # 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):

        print 'auto_import_switch'
        db = state.db
        config = state.config

        # 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 InputError("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:
            print '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()
        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...
            #
            # 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 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()
        for port_name in ports:
            port_id = None
            port_mode = s.port_get_mode(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_vlan = s.port_get_access_vlan(port_name)
                port_id = db.create_port(switch_id, port_name, port_vlan, port_vlan)
                # Nothing further needed
            elif port_mode == 'trunk':
                # 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 = s.port_get_trunk_vlan_list(port_name)
                if port_vlans == 'ALL':
                    port_vlans = (state.default_vlan_id,) # easy for our purposes
                port_id = db.create_port(switch_id, port_name,
                                         port_vlans[0], port_vlans[0])
                # Append to a list of trunk ports that we will need to
                # modify once we're done
                trunk_ports.append(port_id)

        # Done with this switch in the main loop here
        s.switch_disconnect()
        del s

        # Now on each trunk port we need to add all the VLANs across
        # our system, so it's ready to go. Each time we call this, it
        # will connect, disconnect, save config etc.
        for port_id in trunk_ports:
            self.set_port_mode(state, port_id, 'trunk')

        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