Mstsc-Script-Hook/Plugin/ProcessLauncher.cpp

#include "ProcessLauncher.h"
#include "Logging.h"

#include <algorithm>
#include <cstdint>
#include <cwctype>
#include <iomanip>
#include <sstream>
#include <vector>

namespace
{
    std::wstring QuoteWindowsArgument(const std::wstring& arg)
    {
        if (arg.empty())
        {
            return L"\"\"";
        }

        if (arg.find_first_of(L" \t\n\v\"") == std::wstring::npos)
        {
            return arg;
        }

        std::wstring result = L"\"";
        size_t backslashes = 0;
        for (wchar_t ch : arg)
        {
            if (ch == L'\\')
            {
                ++backslashes;
                continue;
            }

            if (ch == L'\"')
            {
                result.append(backslashes * 2 + 1, L'\\');
                result.push_back(L'\"');
                backslashes = 0;
                continue;
            }

            result.append(backslashes, L'\\');
            backslashes = 0;
            result.push_back(ch);
        }

        result.append(backslashes * 2, L'\\');
        result.push_back(L'\"');
        return result;
    }

    std::wstring GetExtensionLower(const std::wstring& path)
    {
        const size_t dot = path.find_last_of(L'.');
        const size_t slash = path.find_last_of(L"\\/");
        if (dot == std::wstring::npos || (slash != std::wstring::npos && dot < slash))
        {
            return {};
        }

        std::wstring ext = path.substr(dot);
        std::transform(ext.begin(), ext.end(), ext.begin(), [](wchar_t c) { return static_cast<wchar_t>(std::towlower(c)); });
        return ext;
    }

    std::wstring ParentDirectory(const std::wstring& path)
    {
        const size_t slash = path.find_last_of(L"\\/");
        if (slash == std::wstring::npos)
        {
            return {};
        }
        if (slash == 2 && path.size() >= 3 && path[1] == L':')
        {
            return path.substr(0, 3);
        }
        return path.substr(0, slash);
    }

    std::wstring GetEnvironmentValue(const wchar_t* name)
    {
        DWORD size = GetEnvironmentVariableW(name, nullptr, 0);
        if (size == 0)
        {
            return {};
        }
        std::vector<wchar_t> buffer(size, L'\0');
        if (GetEnvironmentVariableW(name, buffer.data(), size) == 0)
        {
            return {};
        }
        return std::wstring(buffer.data());
    }

    uint64_t HashPath(const std::wstring& path)
    {
        constexpr uint64_t fnvOffset = 14695981039346656037ull;
        constexpr uint64_t fnvPrime = 1099511628211ull;
        uint64_t hash = fnvOffset;
        for (wchar_t ch : path)
        {
            wchar_t lower = static_cast<wchar_t>(std::towlower(ch));
            hash ^= static_cast<uint16_t>(lower);
            hash *= fnvPrime;
        }
        return hash;
    }

    std::wstring MutexNameForScript(const std::wstring& scriptPath)
    {
        std::wstringstream stream;
        stream << L"Local\\Plandent.MstscScriptHook."
               << std::uppercase << std::hex << std::setw(16) << std::setfill(L'0')
               << HashPath(scriptPath);
        return stream.str();
    }

    struct LaunchCommand
    {
        std::wstring application;
        std::wstring commandLine;
    };

    bool BuildLaunchCommand(const hook::HookConfig& config, LaunchCommand& command)
    {
        const std::wstring extension = GetExtensionLower(config.scriptPath);

        if (extension == L".bat" || extension == L".cmd")
        {
            std::wstring comspec = GetEnvironmentValue(L"ComSpec");
            if (comspec.empty())
            {
                std::wstring systemRoot = GetEnvironmentValue(L"SystemRoot");
                if (systemRoot.empty())
                {
                    return false;
                }
                comspec = systemRoot + L"\\System32\\cmd.exe";
            }

            std::wstring inner = QuoteWindowsArgument(config.scriptPath);
            if (!config.arguments.empty())
            {
                inner += L" " + config.arguments;
            }

            command.application = comspec;
            command.commandLine = QuoteWindowsArgument(comspec) + L" /D /S /C \"" + inner + L"\"";
            return true;
        }

        if (extension == L".ps1")
        {
            std::wstring systemRoot = GetEnvironmentValue(L"SystemRoot");
            if (systemRoot.empty())
            {
                return false;
            }

            std::wstring powershell = systemRoot + L"\\System32\\WindowsPowerShell\\v1.0\\powershell.exe";
            command.application = powershell;
            command.commandLine = QuoteWindowsArgument(powershell) +
                L" -NoLogo -NoProfile -NonInteractive -ExecutionPolicy Bypass -File " +
                QuoteWindowsArgument(config.scriptPath);
            if (!config.arguments.empty())
            {
                command.commandLine += L" " + config.arguments;
            }
            return true;
        }

        if (extension == L".exe")
        {
            command.application = config.scriptPath;
            command.commandLine = QuoteWindowsArgument(config.scriptPath);
            if (!config.arguments.empty())
            {
                command.commandLine += L" " + config.arguments;
            }
            return true;
        }

        return false;
    }
}

namespace hook
{
    ProcessLauncher::~ProcessLauncher()
    {
        if (process_ != nullptr)
        {
            CloseHandle(process_);
            process_ = nullptr;
        }
    }

    void ProcessLauncher::CleanupExitedProcess()
    {
        if (process_ == nullptr)
        {
            return;
        }

        if (WaitForSingleObject(process_, 0) != WAIT_TIMEOUT)
        {
            CloseHandle(process_);
            process_ = nullptr;
            processId_ = 0;
        }
    }

    bool ProcessLauncher::IsRunning()
    {
        CleanupExitedProcess();
        return process_ != nullptr;
    }

    bool ProcessLauncher::Start(const HookConfig& config)
    {
        CleanupExitedProcess();
        if (process_ != nullptr)
        {
            Log(config.enableLogging, L"Script ist für diese MSTSC-Instanz bereits aktiv (PID " + std::to_wstring(processId_) + L").");
            return true;
        }

        if (config.scriptPath.empty())
        {
            Log(config.enableLogging, L"Kein ScriptPath konfiguriert; Start übersprungen.");
            return false;
        }

        DWORD attributes = GetFileAttributesW(config.scriptPath.c_str());
        if (attributes == INVALID_FILE_ATTRIBUTES || (attributes & FILE_ATTRIBUTE_DIRECTORY) != 0)
        {
            Log(config.enableLogging, L"Script wurde nicht gefunden: " + config.scriptPath);
            return false;
        }

        LaunchCommand command;
        if (!BuildLaunchCommand(config, command))
        {
            Log(config.enableLogging, L"Nicht unterstützter Script-/Dateityp: " + config.scriptPath);
            return false;
        }

        std::wstring workingDirectory = config.workingDirectory.empty()
            ? ParentDirectory(config.scriptPath)
            : config.workingDirectory;

        HANDLE duplicateMutex = nullptr;
        if (config.preventDuplicates)
        {
            SECURITY_ATTRIBUTES sa{};
            sa.nLength = sizeof(sa);
            sa.bInheritHandle = TRUE;

            std::wstring mutexName = MutexNameForScript(config.scriptPath);
            duplicateMutex = CreateMutexW(&sa, FALSE, mutexName.c_str());
            if (duplicateMutex == nullptr)
            {
                DWORD error = GetLastError();
                Log(config.enableLogging, L"Mehrfachstart-Schutz konnte nicht erstellt werden: " + FormatWin32Error(error));
                return false;
            }

            if (GetLastError() == ERROR_ALREADY_EXISTS)
            {
                CloseHandle(duplicateMutex);
                Log(config.enableLogging, L"Ein Receiver mit demselben ScriptPath läuft bereits; kein weiterer Start.");
                return true;
            }
        }

        PROCESS_INFORMATION pi{};
        DWORD creationFlags = 0;
        if (config.hidden)
        {
            creationFlags |= CREATE_NO_WINDOW;
        }

        std::vector<wchar_t> mutableCommand(command.commandLine.begin(), command.commandLine.end());
        mutableCommand.push_back(L'\0');

        BOOL created = FALSE;

        if (duplicateMutex != nullptr)
        {
            STARTUPINFOEXW startup{};
            startup.StartupInfo.cb = sizeof(startup);
            if (config.hidden)
            {
                startup.StartupInfo.dwFlags |= STARTF_USESHOWWINDOW;
                startup.StartupInfo.wShowWindow = SW_HIDE;
            }

            SIZE_T attributeSize = 0;
            InitializeProcThreadAttributeList(nullptr, 1, 0, &attributeSize);
            std::vector<BYTE> attributeBuffer(attributeSize);
            startup.lpAttributeList = reinterpret_cast<LPPROC_THREAD_ATTRIBUTE_LIST>(attributeBuffer.data());

            if (!InitializeProcThreadAttributeList(startup.lpAttributeList, 1, 0, &attributeSize))
            {
                DWORD error = GetLastError();
                CloseHandle(duplicateMutex);
                Log(config.enableLogging, L"Prozessattributliste konnte nicht initialisiert werden: " + FormatWin32Error(error));
                return false;
            }

            HANDLE handles[] = { duplicateMutex };
            if (!UpdateProcThreadAttribute(
                startup.lpAttributeList,
                0,
                PROC_THREAD_ATTRIBUTE_HANDLE_LIST,
                handles,
                sizeof(handles),
                nullptr,
                nullptr))
            {
                DWORD error = GetLastError();
                DeleteProcThreadAttributeList(startup.lpAttributeList);
                CloseHandle(duplicateMutex);
                Log(config.enableLogging, L"Handle-Vererbung konnte nicht vorbereitet werden: " + FormatWin32Error(error));
                return false;
            }

            creationFlags |= EXTENDED_STARTUPINFO_PRESENT;
            created = CreateProcessW(
                command.application.c_str(),
                mutableCommand.data(),
                nullptr,
                nullptr,
                TRUE,
                creationFlags,
                nullptr,
                workingDirectory.empty() ? nullptr : workingDirectory.c_str(),
                &startup.StartupInfo,
                &pi);

            DeleteProcThreadAttributeList(startup.lpAttributeList);

            // Der Child-Prozess hat bei erfolgreichem Start eine eigene Referenz auf den Mutex.
            // Dadurch bleibt der Mehrfachstart-Schutz auch bestehen, wenn mstsc.exe später endet.
            CloseHandle(duplicateMutex);
        }
        else
        {
            STARTUPINFOW startup{};
            startup.cb = sizeof(startup);
            if (config.hidden)
            {
                startup.dwFlags |= STARTF_USESHOWWINDOW;
                startup.wShowWindow = SW_HIDE;
            }

            created = CreateProcessW(
                command.application.c_str(),
                mutableCommand.data(),
                nullptr,
                nullptr,
                FALSE,
                creationFlags,
                nullptr,
                workingDirectory.empty() ? nullptr : workingDirectory.c_str(),
                &startup,
                &pi);
        }

        if (!created)
        {
            DWORD error = GetLastError();
            Log(config.enableLogging, L"Scriptstart fehlgeschlagen: " + FormatWin32Error(error));
            return false;
        }

        CloseHandle(pi.hThread);
        process_ = pi.hProcess;
        processId_ = pi.dwProcessId;

        Log(config.enableLogging,
            L"Script gestartet. PID=" + std::to_wstring(processId_) +
            L", Script=" + config.scriptPath +
            (config.arguments.empty() ? L"" : L", Argumente=" + config.arguments));
        return true;
    }

    void ProcessLauncher::Stop(bool loggingEnabled)
    {
        CleanupExitedProcess();
        if (process_ == nullptr)
        {
            return;
        }

        const DWORD pid = processId_;
        if (!TerminateProcess(process_, 0))
        {
            DWORD error = GetLastError();
            Log(loggingEnabled, L"Scriptprozess PID " + std::to_wstring(pid) + L" konnte nicht beendet werden: " + FormatWin32Error(error));
        }
        else
        {
            Log(loggingEnabled, L"Beenden des Scriptprozesses PID " + std::to_wstring(pid) + L" wurde bei RDP-Trennung angefordert.");
        }

        CloseHandle(process_);
        process_ = nullptr;
        processId_ = 0;
    }

    void ProcessLauncher::Detach(bool loggingEnabled)
    {
        CleanupExitedProcess();
        if (process_ == nullptr)
        {
            return;
        }

        Log(loggingEnabled, L"MSTSC beendet Plugin; Script PID " + std::to_wstring(processId_) + L" läuft weiter.");
        CloseHandle(process_);
        process_ = nullptr;
        processId_ = 0;
    }
}