DeployStudio: Verify that the client only has a single volume named Macintosh HD on the boot drive.

dafgu_check_extra_volumes.py

#!/usr/bin/python
# -*- coding: utf-8 -*-
#
# Verify that the client only has a single volume on the boot drive.


import sys
import argparse
import subprocess
import plistlib


def print8(*args):
    print " ".join(unicode(x).encode(u"utf-8") for x in args)


class DiskUtilException(BaseException):
    pass

def diskutil(*args):
    verb = list(args[:-1])
    target = args[-1]
    command = [u"/usr/sbin/diskutil"] + verb + [u"-plist"]
    if target is not None:
        command.append(target)
    try:
        output = subprocess.check_output(command)
        plist = plistlib.readPlistFromString(output)
    except BaseException as e:
        raise DiskUtilException(unicode(e))
    return plist


def main(argv):
    p = argparse.ArgumentParser()
    p.add_argument(u"-v", u"--verbose", action=u"store_true",
                   help=u"Verbose output.")
    args = p.parse_args([x.decode(u"utf-8") for x in argv[1:]])
    
    # Use diskutil to get information about disk0.
    try:
        disk0_plist = diskutil(u"list", u"disk0")
    except DiskUtilException as e:
        print8(u"RuntimeAbortWorkflow: failed to list partitions on disk0: %s" % repr(e))
        return 1
    
    # A single volume named Macintosh HD is the ideal case. Much happiness.
    if disk0_plist[u"VolumesFromDisks"] == [u"Macintosh HD"]:
        print8(u"Found a single volume on disk0 named Macintosh HD")
        return 0
    
    # Bail if there are multiple volumes on disk0.
    if len(disk0_plist[u"VolumesFromDisks"]) > 1:
        print8(u"RuntimeAbortWorkflow: more than one volume on disk0")
        return 1
    
    # Since there are apparently no volumes on disk0, figure out if it's a
    # CoreStorage physical disk.
    try:
        cs_plist = diskutil(u"cs", u"list", None)
    
        # No CoreStorage here, guess the disk is just empty.
        if len(cs_plist) == 0:
            print8(u"RuntimeAbortWorkflow: No CoreStorage volumes found")
            return 1
    
        # Check each LVG.
        for lvg in cs_plist[u"CoreStorageLogicalVolumeGroups"]:
            on_disk0 = False
            # Check the PVs in the LVG.
            for pv in lvg[u"CoreStoragePhysicalVolumes"]:
                pv_plist = diskutil(u"cs", u"info", pv[u"CoreStorageUUID"])
                # Check if the PV is stored on disk0.
                if pv_plist[u"DeviceIdentifier"].startswith(u"disk0"):
                    on_disk0 = True
                    break
            # If disk0 is a member of the LVG, check its volumes.
            if on_disk0:
                # So check each family...
                for fam in lvg[u"CoreStorageLogicalVolumeFamilies"]:
                    # Bail if there are multiple volumes in the LVG.
                    if len(fam[u"CoreStorageLogicalVolumes"]) != 1:
                        print8(u"RuntimeAbortWorkflow: more than one CS volume on disk0")
                        return 1
                    # Make sure the volume is called Macintosh HD.
                    for vol in fam[u"CoreStorageLogicalVolumes"]:
                        vol_plist = diskutil(u"cs", u"info", vol[u"CoreStorageUUID"])
                        if vol_plist[u"VolumeName"] == u"Macintosh HD":
                            print8(u"Found Macintosh HD as CoreStorage volume on disk0")
                        else:
                            print8(u"RuntimeAbortWorkflow: CoreStorage on boot drive should contain a single volume named Macintosh HD")
                            return 1
            else:
                print8(u"Found no volumes on disk0 in LVG %s" % lvg[u"CoreStorageUUID"])
    
    except DiskUtilException as e:
        print8(u"RuntimeAbortWorkflow: failed to list CoreStorage info: %s" % repr(e))
        return 1
    
    
    return 0
    

if __name__ == '__main__':
    sys.exit(main(sys.argv))