dieknow.py

"""DieKnow API.
"""

import ctypes
try:
    from ctypes import wintypes

except ImportError as exc:
    raise OSError("Failed to load Windows ctypes! Ensure you are on a Windows "
                  "platform!") from exc

import os
import sys


if "win" not in sys.platform:
    raise OSError("DieKnow does not support the platform " + sys.platform + \
        ". Only Windows is supported!")


RED = "\033[91m"
RESET = "\033[0m"

md = False
if "-docs" in sys.argv:
    print("Documentation generation enabled")
    md = True

lib_dll_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "dlls", "dieknow.dll")

files = [
    "msvc/msvcp140.dll",
    "msvc/ucrtbase.dll",
    "msvc/vcruntime140.dll",
    "dieknow/dieknow.dll"
]

missing_dlls = False
force_clean_dlls = "/clean" in sys.argv

if force_clean_dlls:
    print("All DLLs will be cleaned from the repository as DieKnow was run "
          "with \"/clean\".")

for file in files:
    extern = file.replace("msvc/", "").replace("dieknow/", "")
    output = lib_dll_path.replace("dieknow.dll", "") + extern
    if not os.path.exists(output):
        missing_dlls = True

if missing_dlls or force_clean_dlls:
    import urllib.request

    print()
    print("DieKnow dynamic link library DLLs are missing. They will be "
          "automatically downloaded from "
          "https://github.com/eschan145/DieKnow.Binaries/blob/main/. Ensure "
          "sure you have a stable Internet connection for the download to "
          "complete smoothly. The download size is approximately 1.75 MB.")
    print()

    os.makedirs(lib_dll_path.replace("dieknow.dll", ""), exist_ok=True)

    prefix = "https://raw.githubusercontent.com/eschan145/DieKnow.Binaries/main/"

    for file in files:
        extern = file.replace("msvc/", "").replace("dieknow/", "")
        output = lib_dll_path.replace("dieknow.dll", "") + extern

        if force_clean_dlls and os.path.exists(output):
            os.remove(output)

        if os.path.exists(output):
            continue

        print(f"Downloading {extern}...", end=" ")

        def progress(block_num, block_size, total_size):
            """Show download progress."""
            downloaded = block_num * block_size
            percent = min(100, downloaded * 100 / total_size)
            print(f"\rDownloading {extern}... {percent:.2f}%", end="")

        urllib.request.urlretrieve(
            prefix + file,
            output,
            reporthook=progress
        )
        print()

    print(
        "Downloaded all required DLLs:",
        ", ".join(file.replace("msvc/", "").replace("dieknow/", "") for file in files)
    )

    print()

try:
    lib = ctypes.CDLL(lib_dll_path)
except OSError as exc:
    error_message = (
        "OSError: "
        "Failed to load DLL libraries! Ensure that the library is "
        "not corrupted, it uses the same architecture (x64) as your "
        "machine (as well as its dependencies) and is not missing "
        "dependencies if dynamically linked!\n\nRefer to "
        "https://learn.microsoft.com/en-us/windows/win32/debug/"
        "system-error-codes for more information."
    )
    sys.stderr.write(f"{RED}{error_message}{RESET}\n")
    raise OSError from exc

try:
    lib.get_kill_method.argtypes = None
    lib.get_kill_method.restype = ctypes.c_int

    lib.set_kill_method.argtypes = [ctypes.c_int]
    lib.set_kill_method.restype = None

    lib.validate.argtypes = None
    lib.validate.restype = None
    lib.get_folder_path.argtypes = None
    lib.get_folder_path.restype = ctypes.c_char_p
    lib.start_monitoring.argtypes = [ctypes.c_char_p]
    lib.start_monitoring.restype = None
    lib.freeze.argtypes = None
    lib.freeze.restype = None
    lib.stop_monitoring.restype = None
    lib.start_monitoring.restype = None
    lib.stop_monitoring.argtypes = None
    lib.close_application_by_exe.restype = ctypes.c_bool
    lib.get_killed_count.restype = ctypes.c_int
    lib.get_executables_in_folder.argtypes = [ctypes.c_char_p]
    lib.get_executables_in_folder.restype = ctypes.c_char_p
    lib.is_running.restype = ctypes.c_bool
    lib.is_monitoring.restype = ctypes.c_bool
    lib.bsod.restype = ctypes.c_int
    lib.dialog.argtypes = [wintypes.LPCWSTR, wintypes.LPCWSTR, wintypes.UINT]
    lib.dialog.restype = ctypes.c_int
except AttributeError as exc:
    parts = str(exc).split("'")
    function_name = parts[1] if len(parts) > 1 else "unknown"
    raise OSError(f"Function '{function_name}' name or ordinal missing!") from exc


for func_name in dir(lib):
    if not func_name.startswith("_"):
        globals()[func_name] = getattr(lib, func_name)

folder_path = lib.get_folder_path()

del func_name  # pylint: disable=W0631

if "/doc" in sys.argv:
    import doc
    doc.doc(os.path.join(os.path.dirname(__file__), "api.cpp"), lib, markdown=md)

try:
    create_window = lib.create_window
except AttributeError as exc:
    parts = str(exc).split("'")
    function_name = parts[1] if len(parts) > 1 else "unknown"
    raise OSError(f"Function '{function_name}' name or ordinal missing!") from exc

# Aliases
gui = create_window
dir = get_executables_in_folder
start = start_monitoring
stop = stop_monitoring
status = is_running
dstatus = is_monitoring

class Shell:
    """DieKnow Shell commands."""

    def __init__(self):
        exclude = {
            "close_application_by_exe",
            "ctypes",
            "doc",
            "get_executables_in_folder",
            "get_folder_path",
            "get_killed_count",
            "create_window",
            "folder_path",
            "is_running",
            "is_monitoring",
            "lib",
            "lib_dll_path",
            "md",
            "RED",
            "RESET",
            "os",
            "start_monitoring",
            "stop_monitoring",
            "sys",
            "wintypes",
            "Shell",
            "get_kill_method",
            "set_kill_method"
        }
        for var_name, value in globals().items():
            if var_name not in exclude and not var_name.startswith("__"):
                setattr(self, var_name, value)

        self.folder = folder_path


shell = Shell()