Add disk_set and disk_toggle functions, and update valid partition flags

salt/modules/parted.py

@@ -44,6 +44,15 @@
 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'])
+
+VALID_PARTITION_FLAGS = set(['boot', 'root', 'swap', 'hidden', 'raid',
+                             'lvm', 'lba', 'hp-service', 'palo',
+                             'prep', 'msftres', 'bios_grub', 'atvrecv',
+                             'diag', 'legacy_boot', 'msftdata', 'irst',
+                             'esp', 'type'])
+
 
 def __virtual__():
     '''
@@ -641,8 +650,25 @@ def set_(device, minor, flag, state):
     :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: bios_grub, legacy_boot, boot, lba, root, swap, hidden, raid,
-        LVM, PALO, PREP, DIAG
+    Valid flags are:
+      * boot
+      * root
+      * swap
+      * hidden
+      * raid
+      * lvm
+      * lba
+      * hp-service
+      * palo
+      * prep
+      * msftres
+      * bios_grub
+      * atvrecv
+      * diag
+      * legacy_boot
+      * msftdata
+      * irst
+      * esp type
 
     CLI Example:
 
@@ -659,8 +685,7 @@ def set_(device, minor, flag, state):
             'Invalid minor number passed to partition.set'
         )
 
-    if flag not in set(['bios_grub', 'legacy_boot', 'boot', 'lba', 'root',
-                        'swap', 'hidden', 'raid', 'LVM', 'PALO', 'PREP', 'DIAG']):
+    if flag not in VALID_PARTITION_FLAGS:
         raise CommandExecutionError('Invalid flag passed to partition.set')
 
     if state not in set(['on', 'off']):
@@ -691,15 +716,68 @@ def toggle(device, partition, flag):
             'Invalid partition number passed to partition.toggle'
         )
 
-    if flag not in set(['bios_grub', 'legacy_boot', 'boot', 'lba', 'root',
-                        'swap', 'hidden', 'raid', 'LVM', 'PALO', 'PREP', 'DIAG']):
+    if flag not in VALID_PARTITION_FLAGS:
         raise CommandExecutionError('Invalid flag passed to partition.toggle')
 
     cmd = 'parted -m -s {0} toggle {1} {2}'.format(device, partition, flag)
     out = __salt__['cmd.run'](cmd).splitlines()
     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 == []