From 72aab6b540691845ba697bb78b83baaef55be70d Mon Sep 17 00:00:00 2001 From: Preston Landers Date: Sun, 8 Nov 2020 13:34:49 -0600 Subject: [PATCH] Fix problems with command line handling in spawned admin process - uses `subprocess.list2cmdline` Adding working unit tests. To run the tests fully automated with no interaction, UAC prompt has to be disabled at system level so it spawns process as admin without asking. --- .gitignore | 1 + CHANGELOG.md | 10 +++ README.md | 79 +++++++++++++++++++++-- pytest.ini | 3 + pyuac/__init__.py | 118 ++++++++-------------------------- pyuac/__version__.py | 2 +- pyuac/admin.py | 96 ++++++++++++++++++++++++++++ pyuac/decorator.py | 127 +++++++++++++++++++++++++++++++++++++ pyuac/run_function.py | 112 -------------------------------- setup.py | 8 ++- tests/example_usage.py | 51 +++++++++++++++ tests/test_decorator.py | 102 +++++++++++++++++++++++++++++ tests/test_main.py | 42 ++++++++++++ tests/test_run_function.py | 36 ----------- 14 files changed, 540 insertions(+), 247 deletions(-) create mode 100644 pytest.ini create mode 100644 pyuac/admin.py create mode 100644 pyuac/decorator.py delete mode 100644 pyuac/run_function.py create mode 100644 tests/example_usage.py create mode 100644 tests/test_decorator.py create mode 100644 tests/test_main.py delete mode 100644 tests/test_run_function.py diff --git a/.gitignore b/.gitignore index c13e2a8..4a4dc33 100644 --- a/.gitignore +++ b/.gitignore @@ -9,3 +9,4 @@ venv27 .eggs .tox .pytest_cache +example_usage.log diff --git a/CHANGELOG.md b/CHANGELOG.md index 60c5eaf..228e98d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,3 +5,13 @@ - Initial package release based on a cleaned up version of my Gist originally published here: https://gist.github.com/Preston-Landers/267391562bc96959eb41 + + Compared to the original version of the Gist, this new version: + + * Has been updated for modern Python 3 (and 2.7) + + * Correctly handles complex script command lines, e.g. quoted spaces in parameters. + + * Provides a decorator for your script's main function for easier usage. + + * It has some basic tests. diff --git a/README.md b/README.md index f3ba0db..e791770 100644 --- a/README.md +++ b/README.md @@ -5,9 +5,69 @@ This package provides a way to invoke User Access Control (UAC) in Windows from This allows a Python process to re-spawn a new process with Administrator level rights using the UAC prompt. Note that the original process is not elevated; a new process is created. -## Quick Usage +The main purpose of pyuac is to allow command line Python scripts to ensure they are run +as Administrator on Windows. There is no ability to execute only parts of a program +as Administrator - the entire script is re-launched with the same command line. You can +also override the command line used for the admin process. -There are two main functions provided: +## Usage and examples + +There are two basic ways to use this library. Perhaps the simplest way is to decorate your +Python command line script's main function. The other is to directly use the `isUserAdmin` +and `runAsAdmin` functions yourself. The decorator allows you to automatically capture +the output of the Admin process and return that output string to the non-admin parent process. + +See also [tests/example_usage.py](tests/example_usage.py) + +### Decorator + +The decorator is an easy way to ensure your script's main() function will respawn itself +as Admin if necessary. + +#### Decorator usage example + +```python +from pyuac import main_requires_admin + +@main_requires_admin +def main(): + print("Do stuff here that requires being run as an admin.") + # The window will disappear as soon as the program exits! + input("Press enter to close the window. >") + +if __name__ == "__main__": + main() +``` + +#### Capture stdout from admin process + +You can also capture the stdout and stderr of your Admin sub-process if you need to check +it for errors from the non-admin parent. By default, unless you set scan_for_error=False on +the decorator, it will check the last line of both stdout and stderr for the words 'error' +or 'exception', and if it finds those, will raise RuntimeError on the parent non-admin side. + +```python +from pyuac import main_requires_admin + +@main_requires_admin(return_output=True) +def main(): + print("Do stuff here that requires being run as an admin.") + # The window will disappear as soon as the program exits! + input("Press enter to close the window. >") + +if __name__ == "__main__": + rv = main() + if not rv: + print("I must have already been Admin!") + else: + admin_stdout, admin_str, *_ = rv + if "Do stuff" in admin_stdout: + print("It worked.") +``` + +### Direct usage + +There are two main direct usage functions provided: isUserAdmin() This returns a boolean to indicate whether the current user has elevated Administrator status. @@ -16,7 +76,7 @@ This returns a boolean to indicate whether the current user has elevated Adminis Re-launch the current process (or the given command line) as an Administrator. This will trigger the UAC (User Access Control) prompt if necessary. -### Example Usage +#### Direct usage example This shows a typical usage pattern: @@ -25,27 +85,32 @@ import pyuac def main(): print("Do stuff here that requires being run as an admin.") + # The window will disappear as soon as the program exits! + input("Press enter to close the window. >") if __name__ == "__main__": if not pyuac.isUserAdmin(): print("Re-launching as admin!") pyuac.runAsAdmin() - else: - main() + else: + main() # Already an admin here. ``` ## Requirements * This package only supports Windows at the moment. The isUserAdmin function will work under - Linux / Posix, but the runAsAdmin command is currently Windows only. + Linux / Posix, but the runAsAdmin functionality is currently Windows only. * This requires Python 2.7, or Python 3.3 or higher. -* This requires the PyWin32 package to be installed. +* This requires the [PyWin32](https://pypi.org/project/pywin32/) package to be installed. https://pypi.org/project/pywin32/ https://github.com/mhammond/pywin32 +* It also depends on the packages [decorator](https://pypi.org/project/decorator/) +and [tee](https://pypi.org/project/tee/) + ## PyWin32 problems The PyWin32 package is required by this library (pyuac). diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..3d61720 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +filterwarnings = + ignore::DeprecationWarning:pywintypes diff --git a/pyuac/__init__.py b/pyuac/__init__.py index 0eaf1c6..6d38eb4 100644 --- a/pyuac/__init__.py +++ b/pyuac/__init__.py @@ -2,8 +2,8 @@ # -*- 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. +""" +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 @@ -12,102 +12,40 @@ Note that the prompt may simply shows a generic python.exe with "Publisher: Unknown" if the python.exe is not signed. However, the standard python.org binaries are signed. -This is meant to be used something like this:: - - if not pyuac.isUserAdmin(): - return pyuac.runAsAdmin() - - # otherwise carry on doing whatever... - -See also this utility function which runs a function as admin and captures the stdout/stderr: - -run_function_as_admin(my_main_function) - -""" - -import os -import sys -from logging import getLogger +This is meant to be used something like this, where you decorate your command line script's +main function: -log = getLogger('pyuac') +>>> from pyuac import main_requires_admin +>>> @main_requires_admin +... def main(): +... # your script main code here. +... return -def isUserAdmin(): - """Check if the current OS user is an Administrator or root. +>>> if __name__ == "__main__": +... main() - :return: True if the current user is an 'Administrator', otherwise False. - """ - if os.name == 'nt': - import win32security +Alternatively, you can do something like this: - try: - adminSid = win32security.CreateWellKnownSid( - win32security.WinBuiltinAdministratorsSid, None) - rv = win32security.CheckTokenMembership(None, adminSid) - log.info("isUserAdmin - CheckTokenMembership returned: %r", rv) - return rv - except Exception as e: - log.warning("Admin check failed, assuming not an admin.", exc_info=e) - return False - else: - # Check for root on Posix - return os.getuid() == 0 +>>> import pyuac +>>> if __name__ == "__main__": +... if not pyuac.isUserAdmin(): +... return pyuac.runAsAdmin() +... # otherwise carry on doing whatever... +... main() -def runAsAdmin(cmdLine=None, wait=True): - """ - Attempt to relaunch the current script as an admin using the same command line parameters. - - WARNING: this function only works on Windows. Future support for Posix might be possible. - Calling this from other than Windows will raise a RuntimeError. - - :param cmdLine: set to override the command line of the program being launched as admin. - Otherwise it defaults to the current process command line! It must be a list in - [command, arg1, arg2...] format. - - :param wait: Set 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. - - :returns: the sub-process return code, unless wait is False, in which case it returns None. - """ - - if os.name != 'nt': - raise RuntimeError("This function is only implemented on Windows.") - - import win32con - import win32event - import win32process - # noinspection PyUnresolvedReferences - from win32com.shell.shell import ShellExecuteEx - # noinspection PyUnresolvedReferences - from win32com.shell import shellcon - - if cmdLine is None: - cmdLine = [sys.executable] + sys.argv - log.debug("Defaulting to runAsAdmin command line: %r", cmdLine) - elif type(cmdLine) not in (tuple, list): - raise ValueError("cmdLine is not a sequence.") - cmd = '"%s"' % (cmdLine[0],) - # XXX TODO: isn't there a function or something we can call to massage command line params? - params = " ".join(['"%s"' % (x,) for x in cmdLine[1:]]) - showCmd = win32con.SW_SHOWNORMAL - lpVerb = 'runas' # causes UAC elevation prompt. +See also this utility function which runs a function as admin and captures the stdout/stderr: - log.info("Running command %r - %r", cmd, params) +run_function_as_admin(my_main_function) - procInfo = ShellExecuteEx( - nShow=showCmd, - fMask=shellcon.SEE_MASK_NOCLOSEPROCESS, - lpVerb=lpVerb, - lpFile=cmd, - lpParameters=params) +""" - if wait: - procHandle = procInfo['hProcess'] - _ = win32event.WaitForSingleObject(procHandle, win32event.INFINITE) - rc = win32process.GetExitCodeProcess(procHandle) - log.info("Process handle %s returned code %s", procHandle, rc) - else: - rc = None +from pyuac.admin import isUserAdmin, runAsAdmin +from pyuac.decorator import main_requires_admin - return rc +__all__ = [ + 'isUserAdmin', + 'runAsAdmin', + 'main_requires_admin' +] diff --git a/pyuac/__version__.py b/pyuac/__version__.py index 3ed72d6..377d7e4 100644 --- a/pyuac/__version__.py +++ b/pyuac/__version__.py @@ -1,7 +1,7 @@ __title__ = 'pyuac' __description__ = 'Python library for Windows User Access Control (UAC)' __url__ = 'https://github.com/Preston-Landers/pyuac' -__version__ = '0.0.1' +__version__ = '0.0.2' __author__ = 'Preston Landers' __author_email__ = 'planders@utexas.edu' __license__ = 'MIT' diff --git a/pyuac/admin.py b/pyuac/admin.py new file mode 100644 index 0000000..036ff3d --- /dev/null +++ b/pyuac/admin.py @@ -0,0 +1,96 @@ +#!/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 + +""" +Contains the core functions. See package docs in __init__.py +""" + +import os +import sys +from logging import getLogger +from subprocess import list2cmdline + +log = getLogger('pyuac') + + +def isUserAdmin(): + """Check if the current OS user is an Administrator or root. + + :return: True if the current user is an 'Administrator', otherwise False. + """ + if os.name == 'nt': + import win32security + + try: + adminSid = win32security.CreateWellKnownSid( + win32security.WinBuiltinAdministratorsSid, None) + rv = win32security.CheckTokenMembership(None, adminSid) + log.info("isUserAdmin - CheckTokenMembership returned: %r", rv) + return rv + except Exception as e: + log.warning("Admin check failed, assuming not an admin.", exc_info=e) + return False + else: + # Check for root on Posix + return os.getuid() == 0 + + +def runAsAdmin(cmdLine=None, wait=True): + """ + Attempt to relaunch the current script as an admin using the same command line parameters. + + WARNING: this function only works on Windows. Future support for Posix might be possible. + Calling this from other than Windows will raise a RuntimeError. + + :param cmdLine: set to override the command line of the program being launched as admin. + Otherwise it defaults to the current process command line! It must be a list in + [command, arg1, arg2...] format. Note that if you're overriding cmdLine, you normally should + make the first element of the list sys.executable + + :param wait: Set 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. + + :returns: the sub-process return code, unless wait is False, in which case it returns None. + """ + + if os.name != 'nt': + raise RuntimeError("This function is only implemented on Windows.") + + import win32con + import win32event + import win32process + # noinspection PyUnresolvedReferences + from win32com.shell.shell import ShellExecuteEx + # noinspection PyUnresolvedReferences + from win32com.shell import shellcon + + if not cmdLine: + cmdLine = [sys.executable] + sys.argv + log.debug("Defaulting to runAsAdmin command line: %r", cmdLine) + elif type(cmdLine) not in (tuple, list): + raise ValueError("cmdLine is not a sequence.") + + showCmd = win32con.SW_SHOWNORMAL + lpVerb = 'runas' # causes UAC elevation prompt. + + cmd = cmdLine[0] + params = list2cmdline(cmdLine[1:]) + + log.info("Running command %r - %r", cmd, params) + procInfo = ShellExecuteEx( + nShow=showCmd, + fMask=shellcon.SEE_MASK_NOCLOSEPROCESS, + lpVerb=lpVerb, + lpFile=cmd, + lpParameters=params) + + if wait: + procHandle = procInfo['hProcess'] + _ = win32event.WaitForSingleObject(procHandle, win32event.INFINITE) + rc = win32process.GetExitCodeProcess(procHandle) + log.info("Process handle %s returned code %s", procHandle, rc) + else: + rc = None + + return rc diff --git a/pyuac/decorator.py b/pyuac/decorator.py new file mode 100644 index 0000000..9878d8a --- /dev/null +++ b/pyuac/decorator.py @@ -0,0 +1,127 @@ +#!/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 + +""" +See main_requires_admin +""" + +import os +import sys +from logging import getLogger + +from decorator import decorator +from tee import StderrTee, StdoutTee + +from pyuac import isUserAdmin, runAsAdmin + +log = getLogger(__name__) + + +@decorator +def main_requires_admin( + run_function, + cmdLine=None, + return_output=False, + stdout_handle=None, stderr_handle=None, + scan_for_error=('error', 'exception'), + *args, **kwargs +): + """ + A decorator for a Python script 'main' function (i.e., when the file is invoked from the + command line as a script) that ensures the 'main' function is executed as an Admin. + Implements a common usage pattern of this module, which allows for capturing the stdout and + stderr output from the sub-process and a very basic scan for errors. + + There is strong assumption here that when the current Python script is re-executed in the + admin context with the same command line args, it code logic will lead to this same decorated + main function being executed again. + + You can NOT send data from the parent (non-admin) process to the child (Admin) + process. Depending on how the parent process was invoked, the child process + might spawn a Python console window that can be interacted with directly. + + Warning: this does not allow capture of the process return code for error detection. + However, the scan_for_error option will look for the case-ins string "error" on the last line + (only) of output, or 'exception', and raise a RuntimeError with the string if found. + + :param run_function: the function to run + :param args: arguments tuple to pass to run_function when called (optional) + :param kwargs: keyword arguments dict to pass to run_function when called (optional) + :param cmdLine: override the command line arguments for the new Admin process. + Defaults to the current command line (sys.argv)! + :param return_output: return the output to the caller of this function instead + of writing it to stdout_handle. Note: due to the nature of how this works with UAC, + this does NOT return the actual "return value" of run_function - only its + stdout and stderr output (as a 2-tuple of (stderr, stdout). If you set this, the callers + of your decorated function should be prepared for this output. + :param stdout_handle: file handle to write the process stdout output, defaults to sys.stdout + :param stderr_handle: file handle to write the process stderr output, defaults to sys.stderr + :param scan_for_error: scan the LAST line only of stdout and stderr for the listed strings. + Case is ignored. Set to None or False to disable this. If one of the listed strings + is found, a RuntimeError is raised in the parent process. + :return: None unless return_output is set. + If return_output is True, the output of the decorated function is a 2-tuple + of (stdout, stderr) strings. + """ + + # Should we add another function parameter to run the in the "not-admin" case? + + # Generate secure temp path? - path has to be known in spawned process... + stdout_temp_fn = 'pyuac.stdout.tmp.txt' + stderr_temp_fn = 'pyuac.stderr.tmp.txt' + + if stdout_handle is None: + stdout_handle = sys.stdout + if stderr_handle is None: + stderr_handle = sys.stderr + + if isUserAdmin(): + with StdoutTee(stdout_temp_fn, mode="a", buff=1), \ + StderrTee(stderr_temp_fn, mode="a", buff=1): + try: + log.debug("Starting run_function as admin") + rv = run_function(*args, **kwargs) + log.debug("Finished run_function as admin. return val: %r", rv) + return rv + except Exception: + log.error("Error running main function as admin", exc_info=True) + raise + else: + log.debug("I'm not admin, starting runAsAdmin") + runAsAdmin(cmdLine=cmdLine, wait=True) + log.debug("I'm not admin, runAsAdmin has finished. Collecting result.") + + rv = [] + for filename, handle in ( + (stdout_temp_fn, stdout_handle), + (stderr_temp_fn, stderr_handle), + ): + if os.path.exists(filename): + with open(filename, "r") as log_fh: + console_output = log_fh.read() + os.remove(filename) + if os.path.exists(filename): + log.warning("Couldn't delete temporary log file %s", filename) + + if scan_for_error: + lines = str.splitlines(console_output.strip()) + if lines: + last_line = lines[-1].strip() + for error_str in scan_for_error: + if last_line.lower().find(error_str) != -1: + log.info( + "Identified an error line in Admin process log at %s - " + "emitting RuntimeError in parent process.\n%s", + filename, last_line) + raise RuntimeError(last_line) + + if return_output: + # return console_output + rv.append(console_output) + + handle.write(console_output) + handle.flush() + + if return_output and rv: + return rv diff --git a/pyuac/run_function.py b/pyuac/run_function.py deleted file mode 100644 index 5e1c82b..0000000 --- a/pyuac/run_function.py +++ /dev/null @@ -1,112 +0,0 @@ -#!/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 - -""" -See run_function_as_admin - -TODO: this is unfinished and untested -""" - -import sys -import os -import traceback - -from tee import StderrTee, StdoutTee - -from pyuac import isUserAdmin, runAsAdmin - - -def run_function_as_admin( - run_function, run_args=None, run_kwargs=None, - return_output=False, stdout_handle=None, stderr_handle=None, - scan_for_error=True, - stdout_temp_fn=None, stderr_temp_fn=None, -): - """ - Implements a common usage pattern of this module, which allows for capturing the stdout and - stderr output from the sub-process. - - You can NOT otherwise interact with the sub-process (supply input to it) at least through - standard input. - - Warning: this does not allow capture of the process return code for error detection. - However, the scan_for_error option will look for the case-ins string "error" on the last line - (only) of output, and raise a RuntimeError with the string if found. - - @param run_function: the function to run - @param run_args: arguments tuple to pass to run_function when called (optional) - @param run_kwargs: keyword arguments dict to pass to run_function when called (optional) - @param return_output: return the output to the caller of this function instead - of writing it to stdout_handle. Note: due to the nature of how this works with UAC, - this does NOT return the actual "return value" of run_function - only its - stdout and stderr output (as a 2-tuple of (stderr, stdout) - @param stdout_handle: file handle to write the process stdout output, defaults to sys.stdout - @param stderr_handle: file handle to write the process stderr output, defaults to sys.stderr - @param scan_for_error: look at the last line only for the string 'error' and - turn that into a RuntimeError if found. ONLY scans the LAST line of stderr and stdout! - @param stdout_temp_fn: the name of the temporary log file to use (will be deleted) - for standard output stream of the sub-process. If not given, a default is generated - @param stderr_temp_fn: the name of the temporary log file to use (will be deleted) - for standard error stream of the sub-process. If not given, a default is generated - @return: None unless return_output is set. - If return_output is True, the output is a 2-tuple of (stdout, stderr) strings. - """ - - # Should we add another function parameter to run the in the "not-admin" case? - - # TODO: generate secure temp path? - if stdout_temp_fn is None: - stdout_temp_fn = 'pyuac.stdout.tmp.txt' - if stderr_temp_fn is None: - stderr_temp_fn = 'pyuac.stderr.tmp.txt' - - if stdout_handle is None: - stdout_handle = sys.stdout - if stderr_handle is None: - stderr_handle = sys.stderr - - if run_kwargs is None: - run_kwargs = {} - if run_args is None: - run_args = () - - if isUserAdmin(): - with StdoutTee(stdout_temp_fn, mode="w"), StderrTee(stderr_temp_fn, mode="w"): - # noinspection PyBroadException - try: - run_function(*run_args, **run_kwargs) - except: - # TODO: re-raise here? - traceback.print_exc(file=stderr_handle) - else: - runAsAdmin(wait=True) - - rv = [] - for filename, handle in ( - (stdout_temp_fn, stdout_handle), - (stderr_temp_fn, stderr_handle), - ): - if os.path.exists(filename): - with open(filename, "r") as log_fh: - console_output = log_fh.read() - os.remove(filename) - if os.path.exists(filename): - print("Warning, can't delete " + filename) - - if scan_for_error: - lines = str.splitlines(console_output.strip()) - if lines: - last_line = lines[-1].strip() - if last_line.lower().find("error") != -1: - raise RuntimeError(last_line) - - if return_output: - # return console_output - rv.append(console_output) - - handle.write(console_output) - handle.flush() - - if return_output and rv: - return rv diff --git a/setup.py b/setup.py index 2c2aeb1..5b506ac 100644 --- a/setup.py +++ b/setup.py @@ -12,7 +12,7 @@ with open(os.path.join(here, 'pyuac', '__version__.py'), 'r', encoding='utf-8') as fh: exec(fh.read(), about) -install_requires = ['tee'] +install_requires = ['tee', 'decorator'] if "win" in sys.platform: try: @@ -45,5 +45,11 @@ "Programming Language :: Python :: 2.7", "License :: OSI Approved :: MIT License", "Operating System :: Microsoft :: Windows", + "Development Status :: 4 - Beta", + "Environment :: Console", + "Environment :: Win32 (MS Windows)", + "Intended Audience :: Developers", + # "Topic :: System :: Shells", + # "Topic :: Utilities", ], ) diff --git a/tests/example_usage.py b/tests/example_usage.py new file mode 100644 index 0000000..7a0f2aa --- /dev/null +++ b/tests/example_usage.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python + +""" +Example demonstrating usage of pyuac, which demonstrates handled quoted strings as +script command line parameters. Invoke this like: + +> python example_usage.py 3 4 --sum --string1 "Here's \"quotes\" in a string!" +""" + +from pyuac.decorator import main_requires_admin + + +@main_requires_admin +def main(): + """ + Some example code; your program here. + """ + import argparse + + parser = argparse.ArgumentParser(description='Process some integers.') + parser.add_argument( + 'integers', metavar='N', type=int, nargs='+', + help='an integer for the accumulator') + parser.add_argument( + '--sum', dest='accumulate', action='store_const', + const=sum, default=max, + help='sum the integers (default: find the max)') + parser.add_argument( + '--string1', dest='string1', action='store', + default=None, + help='Extra string to emit.') + parser.add_argument( + '--string2', dest='string2', action='store', + default=None, + help='Another extra string to emit.') + + args = parser.parse_args() + print(args.accumulate(args.integers)) + + string1 = args.string1 + string2 = args.string2 + if string1 is not None or string2 is not None: + print("String 1: %r" % (string1,)) + print("String 2: %r" % (string2,)) + + input("\nPress enter to quit. >") + + +if __name__ == '__main__': + # Invoke your decorated main function like normal + rv = main() diff --git a/tests/test_decorator.py b/tests/test_decorator.py new file mode 100644 index 0000000..134641f --- /dev/null +++ b/tests/test_decorator.py @@ -0,0 +1,102 @@ +#!/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 + +""" +Test the main_requires_admin decorator + +Note the test is not fully automatic unless the UAC prompt is disabled at the Windows +system level. Otherwise, someone has to click on the UAC prompts. +""" + +import os +import sys + +from pyuac import isUserAdmin, main_requires_admin + +here = os.path.abspath(__file__) + + +def sample_function_body(arg1, kwarg2='Default'): + print("Hello, world.") + is_admin = isUserAdmin() + print("isUserAdmin: %s" % (is_admin,)) + print("arg1: %s" % (arg1,)) + print("kwarg2: %s" % (kwarg2,)) + return + + +example1_args = ('foobar',) +example1_kwargs = {'kwarg2': 'biz'} + + +@main_requires_admin(return_output=True) +def example1_main(): + sample_function_body(*example1_args, **example1_kwargs) + # input("Press Enter >") + return + + +def test_run_function_with_output(): + expected_output = """Hello, world. +isUserAdmin: True +arg1: foobar +kwarg2: biz +""" + cmdLine = [sys.executable, here, 'example1'] + decorated_sample = main_requires_admin( + sample_function_body, return_output=True, cmdLine=cmdLine) + rv = decorated_sample(*example1_args, **example1_kwargs) + assert rv, "Already ran as admin?" + actual_stdout, actual_stderr, *_ = rv + assert actual_stdout == expected_output + assert actual_stderr == "" + return + + +TMP_FILE_NAME = 'test_decorator.py.tmp' +example2_args = ('happy! day',) +example2_kwargs = {'kwarg2': 'Here are \'quotes\' "embedded"'} + + +def sample_function_write_file(arg1, kwarg2='Default'): + with open(TMP_FILE_NAME, "w") as fh: + data = ["Hello, world."] + is_admin = isUserAdmin() + data.append("isUserAdmin: %s" % (is_admin,)) + data.append("arg1: %s" % (arg1,)) + data.append("kwarg2: %s" % (kwarg2,)) + fh.write("\n".join(data) + "\n") + + +@main_requires_admin +def example2_main(): + sample_function_write_file(*example2_args, **example2_kwargs) + # input("Press Enter >") + return + + +def test_run_function_no_output(): + expected_output = """Hello, world. +isUserAdmin: True +arg1: happy! day +kwarg2: Here are 'quotes' "embedded" +""" + cmdLine = [sys.executable, here, 'example2 has spaces and "quote\'s"'] + decorated_sample = main_requires_admin( + sample_function_body, return_output=False, cmdLine=cmdLine) + rv = decorated_sample(*example2_args, **example2_kwargs) + assert rv is None + with open(TMP_FILE_NAME, 'r') as fh: + actual_output = fh.read() + if os.path.exists(TMP_FILE_NAME): + os.unlink(TMP_FILE_NAME) + assert actual_output == expected_output + return + + +if __name__ == '__main__': + if sys.argv[1] == "example1": + example1_main() + elif sys.argv[1] == 'example2 has spaces and "quote\'s"': + example2_main() diff --git a/tests/test_main.py b/tests/test_main.py new file mode 100644 index 0000000..3a1f704 --- /dev/null +++ b/tests/test_main.py @@ -0,0 +1,42 @@ +#!/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 + +""" +Test the basic functions isUserAdmin and runAsAdmin. + +Note the test is not fully automatic unless the UAC prompt is disabled at the Windows +system level. Otherwise, someone has to click on the UAC prompts. +""" + +from __future__ import print_function + +import os +import sys + +from pyuac import isUserAdmin, runAsAdmin + + +def test_1_not_admin(): + """ + The test should not be launched as admin. + :return: + """ + is_admin = isUserAdmin() + assert not is_admin, "You should launch this test as a regular non-admin user." + + +def test_can_run_as_admin(): + """ + A simple test function; check if we're admin, and if not relaunch + the script as admin. + """ + if not isUserAdmin(): + print("You're not an admin.", os.getpid(), "params: ", sys.argv) + # rc = runAsAdmin(["c:\\Windows\\notepad.exe"]) + rc = runAsAdmin() + else: + print("You are an admin!", os.getpid(), "params: ", sys.argv) + rc = 0 + # input('Press Enter to exit. >') + return rc diff --git a/tests/test_run_function.py b/tests/test_run_function.py deleted file mode 100644 index 11f8a42..0000000 --- a/tests/test_run_function.py +++ /dev/null @@ -1,36 +0,0 @@ -#!/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 - -""" - -""" - -from pyuac import isUserAdmin -from pyuac.run_function import run_function_as_admin - - -def sample_function_body(arg1, kwarg2='Default'): - print("Hello, world.") - is_admin = isUserAdmin() - print("isUserAdmin: %s" % (is_admin,)) - print("arg1: %s" % (arg1,)) - print("kwarg2: %s" % (kwarg2,)) - return - - -def test_run_function(): - expected_output = """Hello, world. -isUserAdmin: True -arg1: foobar -kwarg2: biz -""" - actual_stdout, actual_stderr = run_function_as_admin( - sample_function_body, - ('foobar',), - {'kwarg2': 'biz'}, - return_output=True - ) - assert actual_stdout == expected_output - assert actual_stderr == "" - return