PowerToys/src/modules/Workspaces/WorkspacesLib/PwaHelper.cpp

#include "pch.h"
#include "PwaHelper.h"
#include "WindowUtils.h"

#include <filesystem>

#include <appmodel.h>
#include <shellapi.h>
#include <ShlObj.h>
#include <shobjidl.h>
#include <tlhelp32.h>
#include <wrl.h>
#include <propkey.h>

#include <wil/com.h>

#include <common/logger/logger.h>
#include <common/utils/winapi_error.h>

#include <WorkspacesLib/AppUtils.h>
#include <WorkspacesLib/CommandLineArgsHelper.h>
#include <WorkspacesLib/StringUtils.h>

namespace Utils
{
    namespace NonLocalizable
    {
        const std::wstring EdgeAppIdIdentifier = L"--app-id=";
        const std::wstring ChromeAppIdIdentifier = L"Chrome._crx_";
        const std::wstring ChromeBase = L"Google\\Chrome\\User Data\\Default\\Web Applications";
        const std::wstring EdgeBase = L"Microsoft\\Edge\\User Data\\Default\\Web Applications";
        const std::wstring ChromeDirPrefix = L"_crx_";
        const std::wstring EdgeDirPrefix = L"_crx__";
        const std::wstring IcoExtension = L".ico";
    }

    static const std::wstring& GetLocalAppDataFolder()
    {
        static std::wstring localFolder{};

        if (localFolder.empty())
        {
            wil::unique_cotaskmem_string folderPath;
            HRESULT hres = SHGetKnownFolderPath(FOLDERID_LocalAppData, 0, NULL, &folderPath);
            if (SUCCEEDED(hres))
            {
                localFolder = folderPath.get();
            }
            else
            {
                Logger::error(L"Failed to get the local app data folder path: {}", get_last_error_or_default(hres));
                localFolder = L""; // Ensure it is explicitly set to empty on failure
            }
        }

        return localFolder;
    }

    // Finds all PwaHelper.exe processes with the specified parent process ID
    std::vector<DWORD> FindPwaHelperProcessIds()
    {
        std::vector<DWORD> pwaHelperProcessIds;
        const HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
        if (hSnapshot == INVALID_HANDLE_VALUE)
        {
            Logger::info(L"Invalid handle when creating snapshot for the search for PwaHelper processes");
            return pwaHelperProcessIds;
        }

        PROCESSENTRY32 pe;
        pe.dwSize = sizeof(PROCESSENTRY32);

        if (Process32First(hSnapshot, &pe))
        {
            do
            {
                if (_wcsicmp(pe.szExeFile, L"PwaHelper.exe") == 0)
                {
                    Logger::info(L"Found a PWA process with id {}", pe.th32ProcessID);
                    pwaHelperProcessIds.push_back(pe.th32ProcessID);
                }
            } while (Process32Next(hSnapshot, &pe));
        }

        CloseHandle(hSnapshot);
        return pwaHelperProcessIds;
    }

    PwaHelper::PwaHelper()
    {
        InitChromeAppIds();
        InitEdgeAppIds();
    }

    void PwaHelper::InitAppIds(const std::wstring& browserDataFolder, const std::wstring& browserDirPrefix, const std::function<void(const std::wstring&)>& addingAppIdCallback)
    {
        std::filesystem::path folderPath(GetLocalAppDataFolder());
        folderPath.append(browserDataFolder);
        if (!std::filesystem::exists(folderPath))
        {
            Logger::info(L"Edge base path does not exist: {}", folderPath.wstring());
            return;
        }

        try
        {
            for (const auto& directory : std::filesystem::directory_iterator(folderPath))
            {
                if (!directory.is_directory())
                {
                    continue;
                }

                const std::wstring directoryName = directory.path().filename();
                if (directoryName.find(browserDirPrefix) != 0)
                {
                    continue;
                }

                const std::wstring appId = directoryName.substr(browserDirPrefix.length());
                if (addingAppIdCallback)
                {
                    addingAppIdCallback(appId);
                }

                for (const auto& filename : std::filesystem::directory_iterator(directory))
                {
                    if (!filename.is_directory())
                    {
                        const std::filesystem::path filenameString = filename.path().filename();
                        if (StringUtils::CaseInsensitiveEquals(filenameString.extension(), NonLocalizable::IcoExtension))
                        {
                            const auto stem = filenameString.stem().wstring();
                            m_pwaAppIdsToAppNames.insert({ appId, stem });
                            Logger::info(L"Found an installed Pwa app {} with PwaAppId {}", stem, appId);
                            break;
                        }
                    }
                }
            }
        }
        catch (std::exception& ex)
        {
            Logger::error("Failed to iterate over the directory: {}", ex.what());
        }
    }

    void PwaHelper::InitEdgeAppIds()
    {
        if (!m_edgeAppIds.empty())
        {
            // already initialized
            return;
        }

        CommandLineArgsHelper commandLineArgsHelper{};

        const auto pwaHelperProcessIds = FindPwaHelperProcessIds();
        Logger::info(L"Found {} edge Pwa helper processes", pwaHelperProcessIds.size());

        for (const auto subProcessID : pwaHelperProcessIds)
        {
            std::wstring aumidID = GetAUMIDFromProcessId(subProcessID);
            std::wstring commandLineArg = commandLineArgsHelper.GetCommandLineArgs(subProcessID);
            std::wstring appId = GetAppIdFromCommandLineArgs(commandLineArg);

            m_edgeAppIds.insert({ aumidID, appId });
            Logger::info(L"Found an edge Pwa helper process with AumidID {} and PwaAppId {}", aumidID, appId);
        }

        InitAppIds(NonLocalizable::EdgeBase, NonLocalizable::EdgeDirPrefix, [&](const std::wstring&) {});
    }

    void PwaHelper::InitChromeAppIds()
    {
        if (!m_chromeAppIds.empty())
        {
            // already initialized
            return;
        }

        InitAppIds(NonLocalizable::ChromeBase, NonLocalizable::ChromeDirPrefix, [&](const std::wstring& appId) {
            m_chromeAppIds.push_back(appId);
        });
    }

    std::optional<std::wstring> PwaHelper::GetEdgeAppId(const std::wstring& windowAumid) const
    {
        const auto pwaIndex = m_edgeAppIds.find(windowAumid);
        if (pwaIndex != m_edgeAppIds.end())
        {
            return pwaIndex->second;
        }

        return std::nullopt;
    }

    std::optional<std::wstring> PwaHelper::GetChromeAppId(const std::wstring& windowAumid) const
    {
        const auto appIdIndexStart = windowAumid.find(NonLocalizable::ChromeAppIdIdentifier);
        if (appIdIndexStart != std::wstring::npos)
        {
            std::wstring windowAumidSub = windowAumid.substr(appIdIndexStart + NonLocalizable::ChromeAppIdIdentifier.size());
            const auto appIdIndexEnd = windowAumidSub.find(L" ");
            if (appIdIndexEnd != std::wstring::npos)
            {
                windowAumidSub = windowAumidSub.substr(0, appIdIndexEnd);
            }

            const std::wstring windowAumidBegin = windowAumidSub.substr(0, 10);
            for (const auto chromeAppId : m_chromeAppIds)
            {
                if (chromeAppId.find(windowAumidBegin) == 0)
                {
                    return chromeAppId;
                }
            }
        }

        return std::nullopt;
    }

    std::wstring PwaHelper::SearchPwaName(const std::wstring& pwaAppId, const std::wstring& windowAumid) const
    {
        const auto index = m_pwaAppIdsToAppNames.find(pwaAppId);
        if (index != m_pwaAppIdsToAppNames.end())
        {
            return index->second;
        }

        std::wstring nameFromAumid{ windowAumid };
        const std::size_t delimiterPos = nameFromAumid.find(L"-");
        if (delimiterPos != std::string::npos)
        {
            return nameFromAumid.substr(0, delimiterPos);
        }

        return nameFromAumid;
    }

    std::wstring PwaHelper::GetAppIdFromCommandLineArgs(const std::wstring& commandLineArgs) const
    {
        auto result = commandLineArgs;

        // remove the prefix
        if (result.find(NonLocalizable::EdgeAppIdIdentifier) == 0)
        {
            result.erase(0, NonLocalizable::EdgeAppIdIdentifier.length());
        }

        // remove the suffix
        auto appIdIndexEnd = result.find(L" ");
        if (appIdIndexEnd != std::wstring::npos)
        {
            result = result.substr(0, appIdIndexEnd);
        }

        return result;
    }
}