# A module to expose various thread/process/job related structures and
# methods from kernel32
#
# The MIT License
#
# Copyright (c) 2006 the Mozilla Foundation <http://www.mozilla.org>
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.

from ctypes import c_void_p, POINTER, sizeof, Structure, windll, WinError, WINFUNCTYPE
from ctypes.wintypes import BOOL, BYTE, DWORD, HANDLE, LPCWSTR, LPWSTR, UINT, WORD

LPVOID = c_void_p
LPBYTE = POINTER(BYTE)
LPDWORD = POINTER(DWORD)


def ErrCheckBool(result, func, args):
    """errcheck function for Windows functions that return a BOOL True
    on success"""
    if not result:
        raise WinError()
    return args


# CloseHandle()

CloseHandleProto = WINFUNCTYPE(BOOL, HANDLE)
CloseHandle = CloseHandleProto(("CloseHandle", windll.kernel32))
CloseHandle.errcheck = ErrCheckBool

# AutoHANDLE


class AutoHANDLE(HANDLE):
    """Subclass of HANDLE which will call CloseHandle() on deletion."""

    def Close(self):
        if self.value:
            CloseHandle(self)
            self.value = 0

    def __del__(self):
        self.Close()

    def __int__(self):
        return self.value


def ErrCheckHandle(result, func, args):
    """errcheck function for Windows functions that return a HANDLE."""
    if not result:
        raise WinError()
    return AutoHANDLE(result)


# PROCESS_INFORMATION structure


class PROCESS_INFORMATION(Structure):
    _fields_ = [
        ("hProcess", HANDLE),
        ("hThread", HANDLE),
        ("dwProcessID", DWORD),
        ("dwThreadID", DWORD),
    ]

    def __init__(self):
        Structure.__init__(self)

        self.cb = sizeof(self)


LPPROCESS_INFORMATION = POINTER(PROCESS_INFORMATION)

# STARTUPINFO structure


class STARTUPINFO(Structure):
    _fields_ = [
        ("cb", DWORD),
        ("lpReserved", LPWSTR),
        ("lpDesktop", LPWSTR),
        ("lpTitle", LPWSTR),
        ("dwX", DWORD),
        ("dwY", DWORD),
        ("dwXSize", DWORD),
        ("dwYSize", DWORD),
        ("dwXCountChars", DWORD),
        ("dwYCountChars", DWORD),
        ("dwFillAttribute", DWORD),
        ("dwFlags", DWORD),
        ("wShowWindow", WORD),
        ("cbReserved2", WORD),
        ("lpReserved2", LPBYTE),
        ("hStdInput", HANDLE),
        ("hStdOutput", HANDLE),
        ("hStdError", HANDLE),
    ]


LPSTARTUPINFO = POINTER(STARTUPINFO)

STARTF_USESHOWWINDOW = 0x01
STARTF_USESIZE = 0x02
STARTF_USEPOSITION = 0x04
STARTF_USECOUNTCHARS = 0x08
STARTF_USEFILLATTRIBUTE = 0x10
STARTF_RUNFULLSCREEN = 0x20
STARTF_FORCEONFEEDBACK = 0x40
STARTF_FORCEOFFFEEDBACK = 0x80
STARTF_USESTDHANDLES = 0x100

# EnvironmentBlock


class EnvironmentBlock:
    """An object which can be passed as the lpEnv parameter of CreateProcess.
    It is initialized with a dictionary."""

    def __init__(self, dict):
        if not dict:
            self._as_parameter_ = None
        else:
            values = ["%s=%s" % (key, value) for (key, value) in dict.iteritems()]
            values.append("")
            self._as_parameter_ = LPCWSTR("\0".join(values))


# CreateProcess()

CreateProcessProto = WINFUNCTYPE(
    BOOL,  # Return type
    LPCWSTR,  # lpApplicationName
    LPWSTR,  # lpCommandLine
    LPVOID,  # lpProcessAttributes
    LPVOID,  # lpThreadAttributes
    BOOL,  # bInheritHandles
    DWORD,  # dwCreationFlags
    LPVOID,  # lpEnvironment
    LPCWSTR,  # lpCurrentDirectory
    LPSTARTUPINFO,  # lpStartupInfo
    LPPROCESS_INFORMATION,  # lpProcessInformation
)

CreateProcessFlags = (
    (1, "lpApplicationName", None),
    (1, "lpCommandLine"),
    (1, "lpProcessAttributes", None),
    (1, "lpThreadAttributes", None),
    (1, "bInheritHandles", True),
    (1, "dwCreationFlags", 0),
    (1, "lpEnvironment", None),
    (1, "lpCurrentDirectory", None),
    (1, "lpStartupInfo"),
    (2, "lpProcessInformation"),
)


def ErrCheckCreateProcess(result, func, args):
    ErrCheckBool(result, func, args)
    # return a tuple (hProcess, hThread, dwProcessID, dwThreadID)
    pi = args[9]
    return (
        AutoHANDLE(pi.hProcess),
        AutoHANDLE(pi.hThread),
        pi.dwProcessID,
        pi.dwThreadID,
    )


CreateProcess = CreateProcessProto(
    ("CreateProcessW", windll.kernel32), CreateProcessFlags
)
CreateProcess.errcheck = ErrCheckCreateProcess

CREATE_BREAKAWAY_FROM_JOB = 0x01000000
CREATE_DEFAULT_ERROR_MODE = 0x04000000
CREATE_NEW_CONSOLE = 0x00000010
CREATE_NEW_PROCESS_GROUP = 0x00000200
CREATE_NO_WINDOW = 0x08000000
CREATE_SUSPENDED = 0x00000004
CREATE_UNICODE_ENVIRONMENT = 0x00000400
DEBUG_ONLY_THIS_PROCESS = 0x00000002
DEBUG_PROCESS = 0x00000001
DETACHED_PROCESS = 0x00000008

# CreateJobObject()

CreateJobObjectProto = WINFUNCTYPE(
    HANDLE, LPVOID, LPCWSTR  # Return type  # lpJobAttributes  # lpName
)

CreateJobObjectFlags = ((1, "lpJobAttributes", None), (1, "lpName", None))

CreateJobObject = CreateJobObjectProto(
    ("CreateJobObjectW", windll.kernel32), CreateJobObjectFlags
)
CreateJobObject.errcheck = ErrCheckHandle

# AssignProcessToJobObject()

AssignProcessToJobObjectProto = WINFUNCTYPE(
    BOOL, HANDLE, HANDLE  # Return type  # hJob  # hProcess
)
AssignProcessToJobObjectFlags = ((1, "hJob"), (1, "hProcess"))
AssignProcessToJobObject = AssignProcessToJobObjectProto(
    ("AssignProcessToJobObject", windll.kernel32), AssignProcessToJobObjectFlags
)
AssignProcessToJobObject.errcheck = ErrCheckBool

# ResumeThread()


def ErrCheckResumeThread(result, func, args):
    if result == -1:
        raise WinError()

    return args


ResumeThreadProto = WINFUNCTYPE(DWORD, HANDLE)  # Return type  # hThread
ResumeThreadFlags = ((1, "hThread"),)
ResumeThread = ResumeThreadProto(("ResumeThread", windll.kernel32), ResumeThreadFlags)
ResumeThread.errcheck = ErrCheckResumeThread

# TerminateJobObject()

TerminateJobObjectProto = WINFUNCTYPE(
    BOOL, HANDLE, UINT  # Return type  # hJob  # uExitCode
)
TerminateJobObjectFlags = ((1, "hJob"), (1, "uExitCode", 127))
TerminateJobObject = TerminateJobObjectProto(
    ("TerminateJobObject", windll.kernel32), TerminateJobObjectFlags
)
TerminateJobObject.errcheck = ErrCheckBool

# WaitForSingleObject()

WaitForSingleObjectProto = WINFUNCTYPE(
    DWORD,  # Return type
    HANDLE,  # hHandle
    DWORD,  # dwMilliseconds
)
WaitForSingleObjectFlags = ((1, "hHandle"), (1, "dwMilliseconds", -1))
WaitForSingleObject = WaitForSingleObjectProto(
    ("WaitForSingleObject", windll.kernel32), WaitForSingleObjectFlags
)

INFINITE = -1
WAIT_TIMEOUT = 0x0102
WAIT_OBJECT_0 = 0x0
WAIT_ABANDONED = 0x0080

# GetExitCodeProcess()

GetExitCodeProcessProto = WINFUNCTYPE(
    BOOL,  # Return type
    HANDLE,  # hProcess
    LPDWORD,  # lpExitCode
)
GetExitCodeProcessFlags = ((1, "hProcess"), (2, "lpExitCode"))
GetExitCodeProcess = GetExitCodeProcessProto(
    ("GetExitCodeProcess", windll.kernel32), GetExitCodeProcessFlags
)
GetExitCodeProcess.errcheck = ErrCheckBool
