From 336e2210939de8d4c59ed2aabebdbe7da5bb11f6 Mon Sep 17 00:00:00 2001
From: Alberto Planas <aplanas@gmail.com>
Date: Mon, 12 Nov 2018 14:27:13 +0100
Subject: [PATCH] parted: support variable length output for print

The command print from parted have a variable lenght output,
depending on the units requested, or the kind of partition.

This patch add the logic to support the full range of outputs
that parted 3.1, and early versions, are generating.

(cherry picked from commit 5adf859e3f1c78a34886c065d8540947a7718120)
---
 salt/modules/parted.py            | 31 +++++++++++++++--------
 tests/unit/modules/test_parted.py | 41 +++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 10 deletions(-)

diff --git a/salt/modules/parted.py b/salt/modules/parted.py
index b5f50e1e6877..79c7dcd1ad25 100644
--- a/salt/modules/parted.py
+++ b/salt/modules/parted.py
@@ -155,7 +155,7 @@ def list_(device, unit=None):
     for line in out:
         if line in ('BYT;', 'CHS;', 'CYL;'):
             continue
-        cols = line.replace(';', '').split(':')
+        cols = line.rstrip(';').split(':')
         if mode == 'info':
             if 7 <= len(cols) <= 8:
                 ret['info'] = {
@@ -177,15 +177,26 @@ def list_(device, unit=None):
                 raise CommandExecutionError(
                     'Problem encountered while parsing output from parted')
         else:
-            if len(cols) == 7:
-                ret['partitions'][cols[0]] = {
-                    'number': cols[0],
-                    'start': cols[1],
-                    'end': cols[2],
-                    'size': cols[3],
-                    'type': cols[4],
-                    'file system': cols[5],
-                    'flags': cols[6]}
+            # Parted (v3.1) have a variable field list in machine
+            # readable output:
+            #
+            # number:start:end:[size:]([file system:name:flags;]|[free;])
+            #
+            # * If units are in CHS 'size' is not printed.
+            # * If is a logical partition with PED_PARTITION_FREESPACE
+            #   set, the last three fields are replaced with the
+            #   'free' text.
+            #
+            fields = ['number', 'start', 'end']
+            if unit != 'chs':
+                fields.append('size')
+            if cols[-1] == 'free':
+                # Drop the last element from the list
+                cols.pop()
+            else:
+                fields.extend(['file system', 'name', 'flags'])
+            if len(fields) == len(cols):
+                ret['partitions'][cols[0]] = dict(six.moves.zip(fields, cols))
             else:
                 raise CommandExecutionError(
                     'Problem encountered while parsing output from parted')
diff --git a/tests/unit/modules/test_parted.py b/tests/unit/modules/test_parted.py
index 0bf9bd53056d..ea0766107055 100644
--- a/tests/unit/modules/test_parted.py
+++ b/tests/unit/modules/test_parted.py
@@ -116,6 +116,12 @@ def parted_print_output(k):
                 '''1:17.4kB:150MB:150MB:ext3::boot;\n'''
                 '''2:3921GB:4000GB:79.3GB:linux-swap(v1)::;\n'''
             ),
+            "valid chs": (
+                '''CHS;\n'''
+                '''/dev/sda:3133,0,2:scsi:512:512:gpt:AMCC 9650SE-24M DISK:;\n'''
+                '''1:0,0,34:2431,134,43:ext3::boot;\n'''
+                '''2:2431,134,44:2492,80,42:linux-swap(v1)::;\n'''
+            ),
             "valid_legacy": (
                 '''BYT;\n'''
                 '''/dev/sda:4000GB:scsi:512:512:gpt:AMCC 9650SE-24M DISK;\n'''
@@ -262,6 +268,41 @@ def test_list__valid_unit_valid_cmd_output(self):
             }
             self.assertEqual(output, expected)
 
+    def test_list__valid_unit_chs_valid_cmd_output(self):
+        with patch('salt.modules.parted._validate_device', MagicMock()):
+            self.cmdrun_stdout.return_value = self.parted_print_output('valid chs')
+            output = parted.list_('/dev/sda', unit='chs')
+            self.cmdrun_stdout.assert_called_once_with('parted -m -s /dev/sda unit chs print')
+            expected = {
+                'info': {
+                    'logical sector': '512',
+                    'physical sector': '512',
+                    'interface': 'scsi',
+                    'model': 'AMCC 9650SE-24M DISK',
+                    'disk': '/dev/sda',
+                    'disk flags': '',
+                    'partition table': 'gpt',
+                    'size': '3133,0,2'
+                },
+                'partitions': {
+                    '1': {
+                        'end': '2431,134,43',
+                        'number': '1',
+                        'start': '0,0,34',
+                        'file system': 'ext3',
+                        'flags': 'boot',
+                        'name': ''},
+                    '2': {
+                        'end': '2492,80,42',
+                        'number': '2',
+                        'start': '2431,134,44',
+                        'file system': 'linux-swap(v1)',
+                        'flags': '',
+                        'name': ''}
+                }
+            }
+            self.assertEqual(output, expected)
+
     def test_list__valid_legacy_cmd_output(self):
         with patch('salt.modules.parted._validate_device', MagicMock()):
             self.cmdrun_stdout.return_value = self.parted_print_output('valid_legacy')