1# Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4"""A module to provide interface to gpt information.
5
6gpt stands for GUID partition table, it is a data structure describing
7partitions present of a storage device. cgpt is a utility which allows to read
8and modify gpt. This module parses cgpt output to create a dictionary
9including information about all defined partitions including their properties.
10It also allows to modify partition properties as required.
11"""
12
13
14class CgptError(Exception):
15    """Cgpt-specific exception."""
16    pass
17
18
19class CgptHandler(object):
20    """Object representing one or more gpts present in the system.
21
22    Attributes:
23      os_if: an instance of OSInterface, initialized by the caller.
24      devices: a dictionary keyed by the storage device names (as in
25               /dev/sda), the contents are dictionaries of cgpt information,
26               where keys are partiton names, and contents are in turn
27               dictionaries of partition properties, something like the below
28               (compressed for brevity):
29      {'/dev/sda': {
30        'OEM': {'partition': 8, 'Type': 'Linux data', 'UUID': 'xxx'},
31        'ROOT-A': {'partition': 3, 'Type': 'ChromeOS rootfs', 'UUID': 'xyz'},
32        'ROOT-C': {'partition': 7, 'Type': 'ChromeOS rootfs', 'UUID': 'xzz'},
33        'ROOT-B': {'partition': 5, 'Type': 'ChromeOS rootfs', 'UUID': 'aaa'},
34        ...
35        }
36     }
37
38    """
39
40    # This dictionary maps gpt attributes the user can modify into the cgpt
41    # utility command line options.
42    ATTR_TO_COMMAND = {'priority': 'P', 'tries': 'T', 'successful': 'S'}
43
44    def __init__(self, os_if):
45        self.os_if = os_if
46        self.devices = {}
47
48    def read_device_info(self, dev_name):
49        """Get device information from cgpt and parse it into a dictionary.
50
51        Inputs:
52          dev_name: a string the Linux storage device name, (i.e. '/dev/sda')
53        """
54
55        device_dump = self.os_if.run_shell_command_get_output(
56                'cgpt show %s' % dev_name)
57        label = None
58        label_data = {}
59        device_data = {}
60        for line in [x.strip() for x in device_dump]:
61            if 'Label:' in line:
62                if label and label not in device_data:
63                    device_data[label] = label_data
64                _, _, partition, _, label = line.split()
65                label = line.split('Label:')[1].strip('" ')
66                label_data = {'partition': int(partition)}
67                continue
68            if ':' in line:
69                name, value = line.strip().split(':')
70                if name != 'Attr':
71                    label_data[name] = value.strip()
72                    continue
73                # Attributes are split around '=', each attribute becomes a
74                # separate partition property.
75                attrs = value.strip().split()
76                for attr in attrs:
77                    name, value = attr.split('=')
78                    label_data[name] = int(value)
79        if label_data:
80            device_data[label] = label_data
81
82        self.devices[dev_name] = device_data
83
84    def get_partition(self, device, partition_name):
85        """Retrieve a dictionary representing a partition on a device.
86
87        Inputs:
88          device: a string, the Linux device name
89          partition_name: a string, the partition name as reported by cgpt.
90
91        Raises:
92          CgptError in case the device or partiton on that device are not
93          known.
94        """
95
96        try:
97            result = self.devices[device][partition_name]
98        except KeyError:
99            raise CgptError('could not retrieve partiton %s of device %s' %
100                            (partition_name, device))
101        return result
102
103    def set_partition(self, device, partition_name, partition_value):
104        """Set partition properties.
105
106        Inputs:
107          device: a string, the Linux device name
108          partition_name: a string, the partition name as reported by cgpt.
109          partiton_value: a dictionary, where keys are strings, names of the
110                  properties which need to be modified, and values are the
111                  values to set the properties to. The only properties which
112                  can be modified are those which are keys of ATTR_TO_COMMAND
113                  defined above.
114        Raises:
115          CgptError in case a property name is not known or not supposed to
116              be modified.
117        """
118
119        current = self.get_partition(device, partition_name)
120        options = []
121        for prop, value in partition_value.iteritems():
122            try:
123                if value == current[prop]:
124                    continue
125                options.append('-%s %d' % (self.ATTR_TO_COMMAND[prop], value))
126            except KeyError:
127                raise CgptError("unknown or immutable property '%s'" % prop)
128
129        if not options:
130            return
131
132        cgpt_add_cmd = 'cgpt add -i %d %s %s' % (current['partition'],
133                                                 ' '.join(options), device)
134        self.os_if.run_shell_command(cgpt_add_cmd, modifies_device=True)
135