CXMAN-210: Add command for modifying pxe interface

You can now run things like:

cxmanage config pxe status
cxmanage config pxe eth1
cxmanage config pxe eth0

The command modifies the u-boot environment underneath in order
to perist the setting.

7 files changed, 251 insertions, 2 deletions

diff --git a/cxmanage/commands/config.py b/cxmanage/commands/config.py
index ca80928..3d5b060 100644
--- a/cxmanage/commands/config.py
+++ b/cxmanage/commands/config.py
@@ -30,7 +30,8 @@
 
 from cxmanage import get_tftp, get_nodes, get_node_strings, run_command
 
-from cxmanage_api.ubootenv import UbootEnv, validate_boot_args
+from cxmanage_api.ubootenv import UbootEnv, validate_boot_args, \
+        validate_pxe_interface
 
 
 def config_reset_command(args):
@@ -92,3 +93,48 @@ def config_boot_status_command(args):
         print "Some errors occured during the command.\n"
 
     return len(errors) > 0
+
+
+def config_pxe_command(args):
+    """set the PXE boot interface"""
+    if args.interface == "status":
+        return config_pxe_status_command(args)
+
+    validate_pxe_interface(args.interface)
+
+    tftp = get_tftp(args)
+    nodes = get_nodes(args, tftp)
+
+    if not args.quiet:
+        print "Setting pxe interface..."
+
+    results, errors = run_command(args, nodes, "set_pxe_interface",
+            args.interface)
+
+    if not args.quiet and not errors:
+        print "Command completed successfully.\n"
+
+    return len(errors) > 0
+
+
+def config_pxe_status_command(args):
+    tftp = get_tftp(args)
+    nodes = get_nodes(args, tftp)
+
+    if not args.quiet:
+        print "Getting pxe interface..."
+    results, errors = run_command(args, nodes, "get_pxe_interface")
+
+    # Print results
+    if results:
+        node_strings = get_node_strings(args, results, justify=True)
+        print "PXE interface"
+        for node in nodes:
+            if node in results:
+                print "%s: %s" % (node_strings[node], results[node])
+        print
+
+    if not args.quiet and errors:
+        print "Some errors occured during the command.\n"
+
+    return len(errors) > 0
diff --git a/cxmanage_api/fabric.py b/cxmanage_api/fabric.py
index 61a0860..9f4dd76 100644
--- a/cxmanage_api/fabric.py
+++ b/cxmanage_api/fabric.py
@@ -494,6 +494,41 @@ class Fabric(object):
         """
         return self._run_on_all_nodes(async, "get_boot_order")
 
+    def set_pxe_interface(self, interface, async=False):
+        """Sets the pxe interface on all nodes.
+
+        >>> fabric.set_pxe_interface(interface='eth0')
+
+        :param interface: Inteface for pxe requests
+        :type interface: string
+        :param async: Flag that determines if the command result (dictionary)
+                      is returned or a Command object (can get status, etc.).
+        :type async: boolean
+
+        """
+        self._run_on_all_nodes(async, "set_pxe_interface", interface)
+
+    def get_pxe_interface(self, async=False):
+        """Gets the pxe interface from all nodes.
+
+        >>> fabric.get_pxe_interface()
+        {
+         0: 'eth0',
+         1: 'eth0',
+         2: 'eth0',
+         3: 'eth0'
+        }
+
+        :param async: Flag that determines if the command result (dictionary)
+                      is returned or a Command object (can get status, etc.).
+        :type async: boolean
+
+        :returns: The boot order of each node on this fabric.
+        :rtype: dictionary or `Task <tasks.html>`__
+
+        """
+        return self._run_on_all_nodes(async, "get_pxe_interface")
+
     def get_versions(self, async=False):
         """Gets the version info from all nodes.
 
diff --git a/cxmanage_api/node.py b/cxmanage_api/node.py
index 3f79398..c083ae7 100644
--- a/cxmanage_api/node.py
+++ b/cxmanage_api/node.py
@@ -874,6 +874,41 @@ class Node(object):
         """
         return self.get_ubootenv().get_boot_order()
 
+    def set_pxe_interface(self, interface):
+        """Sets pxe interface for this node.
+
+        >>> node.set_boot_order('eth0')
+
+        :param interface: Interface pass on to the uboot environment.
+        :type boot_args: string
+
+        """
+        fwinfo = self.get_firmware_info()
+        first_part = self._get_partition(fwinfo, "UBOOTENV", "FIRST")
+        active_part = self._get_partition(fwinfo, "UBOOTENV", "ACTIVE")
+
+        # Download active ubootenv, modify, then upload to first partition
+        image = self._download_image(active_part)
+        ubootenv = self.ubootenv(open(image.filename).read())
+        ubootenv.set_pxe_interface(interface)
+        priority = max(int(x.priority, 16) for x in [first_part, active_part])
+
+        filename = temp_file()
+        with open(filename, "w") as f:
+            f.write(ubootenv.get_contents())
+
+        ubootenv_image = self.image(filename, image.type, False, image.daddr,
+                                    image.skip_crc32, image.version)
+        self._upload_image(ubootenv_image, first_part, priority)
+
+    def get_pxe_interface(self):
+        """Returns the current pxe interface for this node.
+
+        >>> node.get_pxe_interface()
+        'eth0'
+        """
+        return self.get_ubootenv().get_pxe_interface()
+
     def get_versions(self):
         """Get version info from this node.
 
diff --git a/cxmanage_api/ubootenv.py b/cxmanage_api/ubootenv.py
index b5b8272..12d550d 100644
--- a/cxmanage_api/ubootenv.py
+++ b/cxmanage_api/ubootenv.py
@@ -211,6 +211,64 @@ class UbootEnv:
         validate_boot_args(boot_args) # sanity check
         return boot_args
 
+
+    def set_pxe_interface(self, interface):
+        """Sets the interfacespecified in the uboot environment.
+
+        >>> uboot.set_pxe_interface('eth0')
+
+        .. note::
+            * Valid Args: eth0 or eth1
+
+        :param interface: The interface to set.
+        :type boot_args: string
+
+        :raises ValueError: If an invalid interface is specified.
+
+        """
+        validate_pxe_interface(interface)
+        if interface == self.get_pxe_interface():
+            return
+
+        commands = []
+        retry = False
+        reset = False
+
+        if interface == "eth0":
+            self.variables["ethprime"] = "xgmac0"
+        elif (interface == "eth1"):
+            self.variables["ethprime"] = "xgmac1"
+        else:
+            raise ValueError("Invalid pxe interface: %s" % interface)
+
+    def get_pxe_interface(self):
+        """Returns a string representation of the pxe interface.
+
+        >>> uboot.get_pxe_interface()
+        'eth0'
+
+        :returns: Boot order for this U-Boot Environment.
+        :rtype: string
+        :raises Exception: If the u-boot environment value is not recognized.
+
+        """
+
+        # This is based on reading the ethprime environment variable, and
+        # translating from xgmacX to ethX. By default ethprime is not set
+        # and eth0 is the assumed default (NOTE: this is brittle)
+
+        if "ethprime" in self.variables:
+            xgmac = self.variables["ethprime"]
+            if xgmac == "xgmac0":
+                return "eth0"
+            elif (xgmac == "xgmac1"):
+                return "eth1"
+            else:
+                raise Exception("Unrecognized value for ethprime")
+        else:
+            return "eth0"
+
+
     def get_contents(self):
         """Returns a raw string representation of the uboot environment.
 
@@ -253,3 +311,9 @@ def validate_boot_args(boot_args):
                     raise ValueError("Invalid boot arg: %s" % arg)
         else:
             raise ValueError("Invalid boot arg: %s" % arg)
+
+
+def validate_pxe_interface(interface):
+    """ Validate pxe interface. Raises a ValueError if the args are invalid."""
+    if not interface in ["eth0", "eth1"]:
+        raise ValueError("Invalid pxe interface: %s" % interface)
diff --git a/cxmanage_test/fabric_test.py b/cxmanage_test/fabric_test.py
index fb234c5..0dd2593 100644
--- a/cxmanage_test/fabric_test.py
+++ b/cxmanage_test/fabric_test.py
@@ -140,6 +140,18 @@ class FabricTest(unittest.TestCase):
         for node in self.nodes:
             self.assertEqual(node.executed, ["get_boot_order"])
 
+    def test_set_pxe_interface(self):
+        """ Test set_pxe_interface command """
+        self.fabric.set_pxe_interface("eth0")
+        for node in self.nodes:
+            self.assertEqual(node.executed, [("set_pxe_interface", "eth0")])
+
+    def test_get_pxe_interface(self):
+        """ Test get_pxe_interface command """
+        self.fabric.get_pxe_interface()
+        for node in self.nodes:
+            self.assertEqual(node.executed, ["get_pxe_interface"])
+
     def test_get_versions(self):
         """ Test get_versions command """
         self.fabric.get_versions()
@@ -405,6 +417,13 @@ class DummyNode(object):
         self.executed.append("get_boot_order")
         return ["disk", "pxe"]
 
+    def set_pxe_interface(self, interface):
+        self.executed.append(("set_pxe_interface", interface))
+
+    def get_pxe_interface(self):
+        self.executed.append("get_pxe_interface")
+        return "eth0"
+
     def get_versions(self):
         self.executed.append("get_versions")
 
diff --git a/cxmanage_test/node_test.py b/cxmanage_test/node_test.py
index f5fc188..4f843a7 100644
--- a/cxmanage_test/node_test.py
+++ b/cxmanage_test/node_test.py
@@ -276,6 +276,50 @@ class NodeTest(unittest.TestCase):
 
             self.assertEqual(result, ["disk", "pxe"])
 
+    def test_set_pxe_interface(self):
+        """ Test node.set_pxe_interface method """
+        for node in self.nodes:
+            node.set_pxe_interface("eth0")
+
+            partitions = node.bmc.partitions
+            ubootenv_partition = partitions[5]
+            unchanged_partitions = [x for x in partitions
+                    if x != ubootenv_partition]
+
+            self.assertEqual(ubootenv_partition.updates, 1)
+            self.assertEqual(ubootenv_partition.retrieves, 1)
+            self.assertEqual(ubootenv_partition.checks, 1)
+            self.assertEqual(ubootenv_partition.activates, 1)
+
+            for partition in unchanged_partitions:
+                self.assertEqual(partition.updates, 0)
+                self.assertEqual(partition.retrieves, 0)
+                self.assertEqual(partition.checks, 0)
+                self.assertEqual(partition.activates, 0)
+
+    def test_get_pxe_interface(self):
+        """ Test node.get_pxe_interface method """
+        for node in self.nodes:
+            result = node.get_pxe_interface()
+
+            partitions = node.bmc.partitions
+            ubootenv_partition = partitions[5]
+            unchanged_partitions = [x for x in partitions
+                    if x != ubootenv_partition]
+
+            self.assertEqual(ubootenv_partition.updates, 0)
+            self.assertEqual(ubootenv_partition.retrieves, 1)
+            self.assertEqual(ubootenv_partition.checks, 0)
+            self.assertEqual(ubootenv_partition.activates, 0)
+
+            for partition in unchanged_partitions:
+                self.assertEqual(partition.updates, 0)
+                self.assertEqual(partition.retrieves, 0)
+                self.assertEqual(partition.checks, 0)
+                self.assertEqual(partition.activates, 0)
+
+            self.assertEqual(result, "eth0")
+
     def test_get_versions(self):
         """ Test node.get_versions method """
         for node in self.nodes:
diff --git a/scripts/cxmanage b/scripts/cxmanage
index 3e9036f..889fe4e 100755
--- a/scripts/cxmanage
+++ b/scripts/cxmanage
@@ -42,7 +42,8 @@ from cxmanage.commands.mc import mcreset_command
 from cxmanage.commands.fw import fwupdate_command, fwinfo_command
 from cxmanage.commands.sensor import sensor_command
 from cxmanage.commands.fabric import ipinfo_command, macaddrs_command
-from cxmanage.commands.config import config_reset_command, config_boot_command
+from cxmanage.commands.config import config_reset_command, config_boot_command, \
+        config_pxe_command
 from cxmanage.commands.info import info_command
 from cxmanage.commands.ipmitool import ipmitool_command
 from cxmanage.commands.ipdiscover import ipdiscover_command
@@ -250,6 +251,11 @@ def build_parser():
             type=lambda x: [] if x == 'none' else x.split(','))
     boot.set_defaults(func=config_boot_command)
 
+    pxe = config_subs.add_parser('pxe',
+            help='set pxe interface')
+    pxe.add_argument('interface', help='pxe interface to use')
+    pxe.set_defaults(func=config_pxe_command)
+
     #info command
     info = subparsers.add_parser('info', help='get host info')
     info.add_argument('info_type', nargs='?',