Create Choices XML for munki

createchoices.py

#!/usr/local/bin/python3
"""Creates the required Choices XML for munki."""
from __future__ import print_function

import argparse
import os
import plistlib
import subprocess
import sys

# plistlib.readPlist() and plistlib.writePlist() deprecated in Python 3.4+
DEPRECATED = (sys.version_info.major == 3 and sys.version_info.minor > 4)


def _readPlist(path):
    """Read a property list."""
    result = dict()

    if os.path.exists(path):
        # plistlib.readPlist() deprecated in Python 3.4+
        if DEPRECATED:
            with open(path, 'rb') as _f:
                result = plistlib.load(_f)
        else:
            result = plistlib.readPlist(path)

    return result


def _readPlistFromString(obj):
    """Read a property list."""
    result = dict()

    if DEPRECATED:
        result = plistlib.loads(obj)
    else:
        result = plistlib.readPlistFromString(obj)

    return result


def _writePlist(path, data):
    """Write a property list to file."""
    if DEPRECATED:
        with open(path, 'wb') as _f:
            plistlib.dump(data, _f)
    else:
        plistlib.writePlist(data, path)


def _writePlistToString(data):
    """Write a property list to string."""
    result = None

    if DEPRECATED:
        _result = plistlib.dumps(data, sort_keys=True, fmt=plistlib.FMT_XML)

        try:
            result = _result.decode('utf-8').strip()
        except Exception:
            raise
    else:
        result = plistlib.writePlistToString(data)

    return result


def choices_generator(f, output=None):
    """Generates the XML using 'installer' inbuilt."""
    result = None

    _cmd = ['/usr/sbin/installer', '-showChoiceChangesXML', '-pkg', f]
    _subprocess = subprocess.Popen(_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    _result, _error = _subprocess.communicate()

    if _subprocess.returncode == 0:
        _choices = _readPlistFromString(_result)

        if _choices:
            result = list()

            for _choice in _choices:
                if _choice.get('choiceAttribute', None) == 'selected':
                    result.append(_choice)
    else:
        try:
            _error = _error.decode('utf-8')
        except Exception:
            pass

        print(_error, file=sys.stderr)
        sys.exit(_subprocess.returncode)

    if result:
        if output:
            _writePlist(path=output, data=result)
        else:
            print(_writePlistToString(data=result))


def main():
    parser = argparse.ArgumentParser()

    parser.add_argument('-p', '--package',
                        type=str,
                        dest='pkg_path',
                        metavar='<pkg path>',
                        help='Package path to generate munki compatible choices XML.',
                        required=True)

    parser.add_argument('-o', '--output',
                        type=str,
                        dest='output',
                        metavar='<plist path>',
                        help='Path to write Property List XML out to.',
                        required=False)

    args = parser.parse_args()
    args.pkg_path = os.path.expanduser(os.path.expandvars(args.pkg_path))

    if args.output:
        args.output = os.path.expanduser(os.path.expandvars(args.output))

    if not os.path.exists(args.pkg_path):
        print('Specified file path not found.', file=sys.stderr)
        sys.exit(1)

    choices_generator(f=args.pkg_path, output=args.output)


if __name__ == '__main__':
    main()