pyuac - elevate a Python process with UAC on Windows

pyuac.py

#!/usr/bin/env python
# -*- coding: utf-8; mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vim: fileencoding=utf-8 tabstop=4 expandtab shiftwidth=4

"""User Access Control for Microsoft Windows Vista and higher.  This is
only for the Windows platform.

This will relaunch either the current script - with all the same command
line parameters - or else you can provide a different script/program to
run.  If the current user doesn't normally have admin rights, he'll be
prompted for an admin password. Otherwise he just gets the UAC prompt.

Note that the prompt may simply shows a generic python.exe with "Publisher:
Unknown" if the python.exe is not signed.

This is meant to be used something like this::

    if not pyuac.is_user_admin():
        return pyuac.run_as_admin()

    # otherwise carry on doing whatever...

See L{run_as_admin} for the main interface.

"""

from __future__ import print_function
import logging
import os
import sys
import traceback
import types


def is_user_admin():
    """@return: True if the current user is an 'Admin' whatever that
    means (root on Unix), otherwise False.

    Warning: The inner function fails unless you have Windows XP SP2 or
    higher. The failure causes a traceback to be printed and this
    function to return False.
    """

    if os.name == 'nt':
        import ctypes
        # WARNING: requires Windows XP SP2 or higher!
        try:
            return ctypes.windll.shell32.IsUserAnAdmin()
        except:
            traceback.print_exc()
            logging.error("Admin check failed.  Assuming not an admin.")
            return False
    else:
        # Check for root on Posix
        return os.getuid() == 0


def run_as_admin(cmd_line=None, wait=True, pipe_output=False, cmd_dir=None):
    """Attempt to relaunch the current script as an admin using the same
    command line parameters.  Pass cmd_line in to override and set a new
    command.  It must be a list of [command, arg1, arg2...] format.

    Set wait to False to avoid waiting for the sub-process to finish. You
    will not be able to fetch the exit code of the process if wait is
    False.

    Set pipe_output to True to capture the output of the command into a
    temporary file, and return the output as the return value.  This assumes
    piping from the command to the tempfile works like in a Command prompt
    '> file'.  This will override wait, and set it to True.

    Returns the sub-process return code, unless wait is False in which
    case it returns None.

    @WARNING: this function only works on Windows.
    """

    if os.name != 'nt':
        raise RuntimeError("This function is only implemented on Windows.")

    # import win32api
    import win32con
    import win32event
    import win32process
    from win32com.shell.shell import ShellExecuteEx
    from win32com.shell import shellcon

    python_exe = sys.executable

    if cmd_line is None:
        cmd_line = [python_exe] + sys.argv
    elif sys.version_info >= (3, 0) and type(cmd_line) not in (tuple, list):
            raise ValueError("cmd_line is not a sequence.")
    elif sys.version_info < (3, 0) and type(cmd_line) not in (types.TupleType, types.ListType):
            raise ValueError("cmd_line is not a sequence.")
    cmd = '"%s"' % (cmd_line[0], )

    # Condition commands to run as a script
    if cmd.strip('"').lower() == 'cmd':
        cmd_line = list(cmd_line)
        if not cmd_line[1].startswith('/C'):
            cmd_line[1] = '/C %s' % cmd_line[1]
    elif cmd.strip('"').lower() == 'powershell':
        cmd_line = list(cmd_line)
        if not cmd_line[1].startswith('-Command'):
            cmd_line[1] = '-Command "& {%s}"' % cmd_line[1]
    else:
        pass

    # Pipe to files
    if pipe_output:
        import tempfile
        tf = tempfile.NamedTemporaryFile(delete=False)
        if cmd.strip('"').lower() == 'cmd':
            cmd_line[1] += ' > %s' % tf.name
        elif cmd.strip('"').lower() == 'powershell':
            cmd_line[1] = cmd_line[1].replace('}"', ' | Out-File -Encoding UTF8 -filepath %s}"' % tf.name)
        else:
            logging.error("Unknown command for piping output.  Reverting to non-piped output.")
            pipe_output = False
        tf_name = tf.name
        tf.close()

        if pipe_output:
            if not wait:
                wait = True
                logging.warning("Overriding 'wait' to allow piped output")
        else:
            os.remove(tf_name)

    # XXX TODO: isn't there a function or something we can call to massage command line params?
    params = ' '.join(['%s' % (x,) for x in cmd_line[1:]])
    if not cmd_dir:
        cmd_dir = ''
    show_cmd = win32con.SW_SHOWNORMAL
    lpVerb = 'runas'  # causes UAC elevation prompt.

    logging.info("Running ", cmd, params)

    # ShellExecute() doesn't seem to allow us to fetch the PID or handle
    # of the process, so we can't get anything useful from it. Therefore
    # the more complex ShellExecuteEx() must be used.

    # procHandle = win32api.ShellExecute(0, lpVerb, cmd, params, cmd_dir, show_cmd)

    procInfo = ShellExecuteEx(nShow=show_cmd,
                              fMask=shellcon.SEE_MASK_NOCLOSEPROCESS,
                              lpVerb=lpVerb,
                              lpFile=cmd,
                              lpParameters=params,
                              lpDirectory=cmd_dir)

    if wait:
        procHandle = procInfo['hProcess']
        obj = win32event.WaitForSingleObject(procHandle, win32event.INFINITE)
        if pipe_output:
            with open(tf_name, 'rb') as tf:
                rc = tf.read()
            if cmd.strip('"').lower() == 'cmd':
                rc = rc.decode('utf-8')
            elif cmd.strip('"').lower() == 'powershell':
                rc = rc.decode('utf-8-sig')  # Remove the BOM
            else:
                rc = rc.decode('utf-8')
            rc = rc.replace('\r', '').strip()
            os.remove(tf_name)
        else:
            rc = win32process.GetExitCodeProcess(procHandle)
    else:
        rc = False

    return rc


def test():
    """A simple test function; check if we're admin, and if not relaunch
    the script as admin.""",
    rc = False
    if not is_user_admin():
        print("You're not an admin.", os.getpid(), "params: ", sys.argv)
        rc = run_as_admin(sys.argv[1:])
    else:
        print("You are an admin!", os.getpid(), "params: ", sys.argv)
        rc = False
    input('Press Enter to exit.')
    return rc


if __name__ == "__main__":
    res = test()
    sys.exit(res)

This fork makes the following changes:

• Updates for PEP8
• Adds some functionality for returning values from the elevated script to the master script via writing a temp file.
  • Currently supports Windows cmd and powershell as there needs to be some formatting goodies to make sure the output is dumped to and read from the temporary file appropriately.

Tested mainly on Python3, though ugly Python2 syntax was kept for possible simple modifications to revert back to the dark side.