// Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.

#include "precomp.h"


// constants

#define WU_S_REBOOT_REQUIRED        0x00240005L
#define WU_S_ALREADY_INSTALLED      0x00240006L


// function definitions
static HRESULT EnsureWUServiceEnabled(
    __in BOOL fStopWusaService,
    __out SC_HANDLE* pschWu,
    __out BOOL* pfPreviouslyDisabled
    );
static HRESULT SetServiceStartType(
    __in SC_HANDLE sch,
    __in DWORD stratType
    );
static HRESULT StopWUService(
    __in SC_HANDLE schWu
    );


extern "C" HRESULT MsuEngineParsePackageFromXml(
    __in IXMLDOMNode* pixnMsuPackage,
    __in BURN_PACKAGE* pPackage
    )
{
    HRESULT hr = S_OK;

    // @DetectCondition
    hr = XmlGetAttributeEx(pixnMsuPackage, L"DetectCondition", &pPackage->Msu.sczDetectCondition);
    ExitOnRequiredXmlQueryFailure(hr, "Failed to get @DetectCondition.");

    // Uninstalling MSU packages isn't supported because newer OS's don't allow silent uninstallation.
    pPackage->fPermanent = TRUE;

LExit:
    return hr;
}

extern "C" void MsuEnginePackageUninitialize(
    __in BURN_PACKAGE* pPackage
    )
{
    ReleaseNullStr(pPackage->Msu.sczDetectCondition);
}

extern "C" HRESULT MsuEngineDetectPackage(
    __in BURN_PACKAGE* pPackage,
    __in BURN_REGISTRATION* pRegistration,
    __in BURN_VARIABLES* pVariables
    )
{
    HRESULT hr = S_OK;
    BOOL fDetected = FALSE;

    // evaluate detect condition
    if (pPackage->Msu.sczDetectCondition && *pPackage->Msu.sczDetectCondition)
    {
        hr = ConditionEvaluate(pVariables, pPackage->Msu.sczDetectCondition, &fDetected);
        ExitOnFailure(hr, "Failed to evaluate MSU package detect condition.");
    }

    // update detect state
    pPackage->currentState = fDetected ? BOOTSTRAPPER_PACKAGE_STATE_PRESENT : BOOTSTRAPPER_PACKAGE_STATE_ABSENT;

    if (pPackage->fCanAffectRegistration)
    {
        pPackage->installRegistrationState = BOOTSTRAPPER_PACKAGE_STATE_ABSENT < pPackage->currentState ? BURN_PACKAGE_REGISTRATION_STATE_PRESENT : BURN_PACKAGE_REGISTRATION_STATE_ABSENT;
    }

    hr = DependencyDetectChainPackage(pPackage, pRegistration);
    ExitOnFailure(hr, "Failed to detect dependencies for MSU package.");

LExit:
    return hr;
}

//
// PlanCalculate - calculates the execute and rollback state for the requested package state.
//
extern "C" HRESULT MsuEnginePlanCalculatePackage(
    __in BURN_PACKAGE* pPackage
    )
{
    HRESULT hr = S_OK;
    BOOTSTRAPPER_ACTION_STATE execute = BOOTSTRAPPER_ACTION_STATE_NONE;
    BOOTSTRAPPER_ACTION_STATE rollback = BOOTSTRAPPER_ACTION_STATE_NONE;

    // execute action
    switch (pPackage->currentState)
    {
    case BOOTSTRAPPER_PACKAGE_STATE_PRESENT:
        execute = BOOTSTRAPPER_ACTION_STATE_NONE;
        break;

    case BOOTSTRAPPER_PACKAGE_STATE_ABSENT:
        switch (pPackage->requested)
        {
        case BOOTSTRAPPER_REQUEST_STATE_PRESENT: __fallthrough;
        case BOOTSTRAPPER_REQUEST_STATE_FORCE_PRESENT: __fallthrough;
        case BOOTSTRAPPER_REQUEST_STATE_REPAIR:
            execute = BOOTSTRAPPER_ACTION_STATE_INSTALL;
            break;

        default:
            execute = BOOTSTRAPPER_ACTION_STATE_NONE;
            break;
        }
        break;

    default:
        hr = E_INVALIDARG;
        ExitOnRootFailure(hr, "Invalid package state.");
    }

    // Calculate the rollback action if there is an execute action.
    if (BOOTSTRAPPER_ACTION_STATE_NONE != execute)
    {
        rollback = BOOTSTRAPPER_ACTION_STATE_NONE;
    }

    // return values
    pPackage->execute = execute;
    pPackage->rollback = rollback;

LExit:
    return hr;
}

//
// PlanAdd - adds the calculated execute and rollback actions for the package.
//
extern "C" HRESULT MsuEnginePlanAddPackage(
    __in BURN_PACKAGE* pPackage,
    __in BURN_PLAN* pPlan,
    __in BURN_LOGGING* pLog,
    __in BURN_VARIABLES* pVariables
    )
{
    HRESULT hr = S_OK;
    BURN_EXECUTE_ACTION* pAction = NULL;

    hr = DependencyPlanPackage(NULL, pPackage, pPlan);
    ExitOnFailure(hr, "Failed to plan package dependency actions.");

    // add rollback action
    if (BOOTSTRAPPER_ACTION_STATE_NONE != pPackage->rollback)
    {
        hr = PlanAppendRollbackAction(pPlan, &pAction);
        ExitOnFailure(hr, "Failed to append rollback action.");

        pAction->type = BURN_EXECUTE_ACTION_TYPE_MSU_PACKAGE;
        pAction->msuPackage.pPackage = pPackage;
        pAction->msuPackage.action = pPackage->rollback;

        LoggingSetPackageVariable(pPackage, NULL, TRUE, pLog, pVariables, &pAction->msuPackage.sczLogPath); // ignore errors.

        hr = PlanExecuteCheckpoint(pPlan);
        ExitOnFailure(hr, "Failed to append execute checkpoint.");
    }

    // add execute action
    if (BOOTSTRAPPER_ACTION_STATE_NONE != pPackage->execute)
    {
        hr = PlanAppendExecuteAction(pPlan, &pAction);
        ExitOnFailure(hr, "Failed to append execute action.");

        pAction->type = BURN_EXECUTE_ACTION_TYPE_MSU_PACKAGE;
        pAction->msuPackage.pPackage = pPackage;
        pAction->msuPackage.action = pPackage->execute;

        LoggingSetPackageVariable(pPackage, NULL, FALSE, pLog, pVariables, &pAction->msuPackage.sczLogPath); // ignore errors.
    }

LExit:
    return hr;
}

extern "C" HRESULT MsuEngineExecutePackage(
    __in BURN_EXECUTE_ACTION* pExecuteAction,
    __in BURN_CACHE* pCache,
    __in BURN_VARIABLES* pVariables,
    __in BOOL fRollback,
    __in BOOL fStopWusaService,
    __in PFN_GENERICMESSAGEHANDLER pfnGenericMessageHandler,
    __in LPVOID pvContext,
    __out BOOTSTRAPPER_APPLY_RESTART* pRestart
    )
{
    HRESULT hr = S_OK;
    LPWSTR sczCachedDirectory = NULL;
    LPWSTR sczMsuPath = NULL;
    LPWSTR sczSystemPath = NULL;
    LPWSTR sczWusaPath = NULL;
    LPWSTR sczCommand = NULL;
    SC_HANDLE schWu = NULL;
    BOOL fWuWasDisabled = FALSE;
    STARTUPINFOW si = { };
    PROCESS_INFORMATION pi = { };
    GENERIC_EXECUTE_MESSAGE message = { };
    DWORD dwExitCode = 0;
    BOOL fUseSysNativePath = FALSE;
    BURN_PACKAGE* pPackage = pExecuteAction->msuPackage.pPackage;
    BURN_PAYLOAD* pPackagePayload = pPackage->payloads.rgItems[0].pPayload;

#if !defined(_WIN64)
    hr = ProcWow64(::GetCurrentProcess(), &fUseSysNativePath);
    ExitOnFailure(hr, "Failed to determine WOW64 status.");
#endif

    *pRestart = BOOTSTRAPPER_APPLY_RESTART_NONE;

    // get wusa.exe path
    if (fUseSysNativePath)
    {
        hr = PathSystemWindowsSubdirectory(L"SysNative\\", &sczSystemPath);
        ExitOnFailure(hr, "Failed to append SysNative directory.");
    }
    else
    {
        hr = PathGetSystemDirectory(&sczSystemPath);
        ExitOnFailure(hr, "Failed to find System32 directory.");
    }

    hr = PathConcat(sczSystemPath, L"wusa.exe", &sczWusaPath);
    ExitOnFailure(hr, "Failed to allocate WUSA.exe path.");

    // build command
    switch (pExecuteAction->msuPackage.action)
    {
    case BOOTSTRAPPER_ACTION_STATE_INSTALL:
        // get cached MSU path
        hr = CacheGetCompletedPath(pCache, TRUE, pPackage->sczCacheId, &sczCachedDirectory);
        ExitOnFailure(hr, "Failed to get cached path for package: %ls", pPackage->sczId);

        // Best effort to set the execute package cache folder variable.
        VariableSetString(pVariables, BURN_BUNDLE_EXECUTE_PACKAGE_CACHE_FOLDER, sczCachedDirectory, TRUE, FALSE);

        hr = PathConcatRelativeToFullyQualifiedBase(sczCachedDirectory, pPackagePayload->sczFilePath, &sczMsuPath);
        ExitOnFailure(hr, "Failed to build MSU path.");

        // format command
        hr = StrAllocFormatted(&sczCommand, L"\"%ls\" \"%ls\" /quiet /norestart", sczWusaPath, sczMsuPath);
        ExitOnFailure(hr, "Failed to format MSU install command.");
        break;

    default:
        hr = E_UNEXPECTED;
        ExitOnFailure(hr, "Failed to get action arguments for MSU package.");
    }

    if (pExecuteAction->msuPackage.sczLogPath && *pExecuteAction->msuPackage.sczLogPath)
    {
        hr = StrAllocConcat(&sczCommand, L" /log:", 0);
        ExitOnFailure(hr, "Failed to append log switch to MSU command-line.");

        hr = StrAllocConcat(&sczCommand, pExecuteAction->msuPackage.sczLogPath, 0);
        ExitOnFailure(hr, "Failed to append log path to MSU command-line.");
    }

    LogId(REPORT_STANDARD, MSG_APPLYING_PACKAGE, LoggingRollbackOrExecute(fRollback), pPackage->sczId, LoggingActionStateToString(pExecuteAction->msuPackage.action), sczMsuPath, sczCommand);

    hr = EnsureWUServiceEnabled(fStopWusaService, &schWu, &fWuWasDisabled);
    ExitOnFailure(hr, "Failed to ensure WU service was enabled to install MSU package.");

    hr = ExeEngineRunProcess(pfnGenericMessageHandler, pvContext, pPackage, sczWusaPath, sczCommand, NULL, NULL, &dwExitCode);
    ExitOnFailure(hr, "Failed to run MSU process");

    // We'll normalize the restart required error code from wusa.exe just in case. Most likely
    // that on reboot we'll actually get WU_S_REBOOT_REQUIRED.
    if (HRESULT_FROM_WIN32(ERROR_SUCCESS_REBOOT_REQUIRED) == static_cast<HRESULT>(dwExitCode))
    {
        dwExitCode = ERROR_SUCCESS_REBOOT_REQUIRED;
    }

    // handle exit code
    switch (dwExitCode)
    {
    case S_OK: __fallthrough;
    case S_FALSE: __fallthrough;
    case WU_S_ALREADY_INSTALLED:
        hr = S_OK;
        break;

    case ERROR_SUCCESS_REBOOT_REQUIRED: __fallthrough;
    case WU_S_REBOOT_REQUIRED:
        *pRestart = BOOTSTRAPPER_APPLY_RESTART_REQUIRED;
        hr = S_OK;
        break;

    default:
        hr = static_cast<HRESULT>(dwExitCode);
        break;
    }

LExit:
    ReleaseStr(sczCachedDirectory);
    ReleaseStr(sczMsuPath);
    ReleaseStr(sczSystemPath);
    ReleaseStr(sczWusaPath);
    ReleaseStr(sczCommand);

    ReleaseHandle(pi.hProcess);
    ReleaseHandle(pi.hThread);

    if (fWuWasDisabled)
    {
        SetServiceStartType(schWu, SERVICE_DISABLED);
    }

    // Best effort to clear the execute package cache folder variable.
    VariableSetString(pVariables, BURN_BUNDLE_EXECUTE_PACKAGE_CACHE_FOLDER, NULL, TRUE, FALSE);

    return hr;
}

extern "C" void MsuEngineUpdateInstallRegistrationState(
    __in BURN_EXECUTE_ACTION* pAction,
    __in HRESULT hrExecute
    )
{
    BURN_PACKAGE* pPackage = pAction->msuPackage.pPackage;

    if (FAILED(hrExecute) || !pPackage->fCanAffectRegistration)
    {
        ExitFunction();
    }

    if (BOOTSTRAPPER_ACTION_STATE_UNINSTALL == pAction->msuPackage.action)
    {
        pPackage->installRegistrationState = BURN_PACKAGE_REGISTRATION_STATE_ABSENT;
    }
    else
    {
        pPackage->installRegistrationState = BURN_PACKAGE_REGISTRATION_STATE_PRESENT;
    }

LExit:
    return;
}

static HRESULT EnsureWUServiceEnabled(
    __in BOOL fStopWusaService,
    __out SC_HANDLE* pschWu,
    __out BOOL* pfPreviouslyDisabled
    )
{
    HRESULT hr = S_OK;
    SC_HANDLE schSCM = NULL;
    SC_HANDLE schWu = NULL;
    SERVICE_STATUS serviceStatus = { };
    QUERY_SERVICE_CONFIGW* pConfig = NULL;

    schSCM = ::OpenSCManagerW(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    ExitOnNullWithLastError(schSCM, hr, "Failed to open service control manager.");

    schWu = ::OpenServiceW(schSCM, L"wuauserv", SERVICE_QUERY_CONFIG | SERVICE_CHANGE_CONFIG | SERVICE_QUERY_STATUS | SERVICE_STOP );
    ExitOnNullWithLastError(schWu, hr, "Failed to open WU service.");

    if (!::QueryServiceStatus(schWu, &serviceStatus) )
    {
        ExitWithLastError(hr, "Failed to query status of WU service.");
    }

    // Stop service if requested to.
    if (SERVICE_STOPPED != serviceStatus.dwCurrentState && fStopWusaService)
    {
        hr = StopWUService(schWu);
    }

    // If the service is not running then it might be disabled so let's check.
    if (SERVICE_RUNNING != serviceStatus.dwCurrentState)
    {
        hr = SvcQueryConfig(schWu, &pConfig);
        ExitOnFailure(hr, "Failed to read configuration for WU service.");

        // If WU is disabled, change it to a demand start service (but touch nothing else).
        if (SERVICE_DISABLED == pConfig->dwStartType)
        {
            hr = SetServiceStartType(schWu, SERVICE_DEMAND_START);
            ExitOnFailure(hr, "Failed to mark WU service to start on demand.");

            *pfPreviouslyDisabled = TRUE;
        }
    }

    *pschWu = schWu;
    schWu = NULL;

LExit:
    ReleaseMem(pConfig);
    ReleaseServiceHandle(schWu);
    ReleaseServiceHandle(schSCM);

    return hr;
}

static HRESULT SetServiceStartType(
    __in SC_HANDLE sch,
    __in DWORD startType
    )
{
    HRESULT hr = S_OK;

    if (!::ChangeServiceConfigW(sch, SERVICE_NO_CHANGE, startType, SERVICE_NO_CHANGE, NULL, NULL, NULL, NULL, NULL, NULL, NULL))
    {
        ExitWithLastError(hr, "Failed to set service start type.");
    }

LExit:
    return hr;
}

static HRESULT StopWUService(
    __in SC_HANDLE schWu
    )
{
    HRESULT hr = S_OK;
    SERVICE_STATUS serviceStatus = { };

    if(!::ControlService(schWu, SERVICE_CONTROL_STOP, &serviceStatus))
    {
        ExitWithLastError(hr, "Failed to stop wusa service.");
    }

LExit:
    return hr;
}