parted: add disk_{set, toggle} functions

salt/modules/parted.py

@@ -44,6 +44,9 @@
 VALID_UNITS = set(['s', 'B', 'kB', 'MB', 'MiB', 'GB', 'GiB', 'TB', 'TiB', '%',
                    'cyl', 'chs', 'compact'])
 
+VALID_DISK_FLAGS = set(['cylinder_alignment', 'pmbr_boot',
+                        'implicit_partition_table'])
+
 
 def __virtual__():
     '''
@@ -700,6 +703,60 @@ def toggle(device, partition, flag):
     return out
 
 
+def disk_set(device, flag, state):
+    '''
+    Changes a flag on selected device.
+
+    A flag can be either "on" or "off" (make sure to use proper
+    quoting, see :ref:`YAML Idiosyncrasies
+    <yaml-idiosyncrasies>`). Some or all of these flags will be
+    available, depending on what disk label you are using.
+
+    Valid flags are:
+      * cylinder_alignment
+      * pmbr_boot
+      * implicit_partition_table
+
+    CLI Example:
+
+    .. code-block:: bash
+
+        salt '*' partition.disk_set /dev/sda pmbr_boot '"on"'
+    '''
+    _validate_device(device)
+
+    if flag not in VALID_DISK_FLAGS:
+        raise CommandExecutionError('Invalid flag passed to partition.disk_set')
+
+    if state not in set(['on', 'off']):
+        raise CommandExecutionError('Invalid state passed to partition.disk_set')
+
+    cmd = ['parted', '-m', '-s', device, 'disk_set', flag, state]
+    out = __salt__['cmd.run'](cmd).splitlines()
+    return out
+
+
+def disk_toggle(device, flag):
+    '''
+    Toggle the state of <flag> on <device>. Valid flags are the same
+    as the disk_set command.
+
+    CLI Example:
+
+    .. code-block:: bash
+
+        salt '*' partition.disk_toggle /dev/sda pmbr_boot
+    '''
+    _validate_device(device)
+
+    if flag not in VALID_DISK_FLAGS:
+        raise CommandExecutionError('Invalid flag passed to partition.disk_toggle')
+
+    cmd = ['parted', '-m', '-s', device, 'disk_toggle', flag]
+    out = __salt__['cmd.run'](cmd).splitlines()
+    return out
+
+
 def exists(device=''):
     '''
     Check to see if the partition exists

tests/unit/modules/test_parted.py

@@ -376,3 +376,20 @@ def test_list__valid_unit_valid_legacy_cmd_output(self):
                 }
             }
             self.assertEqual(output, expected)
+
+    def test_disk_set(self):
+        with patch('salt.modules.parted._validate_device', MagicMock()):
+            self.cmdrun.return_value = ''
+            output = parted.disk_set('/dev/sda', 'pmbr_boot', 'on')
+            self.cmdrun.assert_called_once_with(
+                ['parted', '-m', '-s', '/dev/sda', 'disk_set',
+                 'pmbr_boot', 'on'])
+            assert output == []
+
+    def test_disk_toggle(self):
+        with patch('salt.modules.parted._validate_device', MagicMock()):
+            self.cmdrun.return_value = ''
+            output = parted.disk_toggle('/dev/sda', 'pmbr_boot')
+            self.cmdrun.assert_called_once_with(
+                ['parted', '-m', '-s', '/dev/sda', 'disk_toggle', 'pmbr_boot'])
+            assert output == []