// 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"


// Exit macros
#define SrpExitOnLastError(x, s, ...) ExitOnLastErrorSource(DUTIL_SOURCE_SRPUTIL, x, s, __VA_ARGS__)
#define SrpExitOnLastErrorDebugTrace(x, s, ...) ExitOnLastErrorDebugTraceSource(DUTIL_SOURCE_SRPUTIL, x, s, __VA_ARGS__)
#define SrpExitWithLastError(x, s, ...) ExitWithLastErrorSource(DUTIL_SOURCE_SRPUTIL, x, s, __VA_ARGS__)
#define SrpExitOnFailure(x, s, ...) ExitOnFailureSource(DUTIL_SOURCE_SRPUTIL, x, s, __VA_ARGS__)
#define SrpExitOnRootFailure(x, s, ...) ExitOnRootFailureSource(DUTIL_SOURCE_SRPUTIL, x, s, __VA_ARGS__)
#define SrpExitOnFailureDebugTrace(x, s, ...) ExitOnFailureDebugTraceSource(DUTIL_SOURCE_SRPUTIL, x, s, __VA_ARGS__)
#define SrpExitOnNull(p, x, e, s, ...) ExitOnNullSource(DUTIL_SOURCE_SRPUTIL, p, x, e, s, __VA_ARGS__)
#define SrpExitOnNullWithLastError(p, x, s, ...) ExitOnNullWithLastErrorSource(DUTIL_SOURCE_SRPUTIL, p, x, s, __VA_ARGS__)
#define SrpExitOnNullDebugTrace(p, x, e, s, ...)  ExitOnNullDebugTraceSource(DUTIL_SOURCE_SRPUTIL, p, x, e, s, __VA_ARGS__)
#define SrpExitOnInvalidHandleWithLastError(p, x, s, ...) ExitOnInvalidHandleWithLastErrorSource(DUTIL_SOURCE_SRPUTIL, p, x, s, __VA_ARGS__)
#define SrpExitOnWin32Error(e, x, s, ...) ExitOnWin32ErrorSource(DUTIL_SOURCE_SRPUTIL, e, x, s, __VA_ARGS__)
#define SrpExitOnGdipFailure(g, x, s, ...) ExitOnGdipFailureSource(DUTIL_SOURCE_SRPUTIL, g, x, s, __VA_ARGS__)


typedef BOOL (WINAPI *PFN_SETRESTOREPTW)(
    __in PRESTOREPOINTINFOW pRestorePtSpec,
    __out PSTATEMGRSTATUS pSMgrStatus
    );

static PFN_SETRESTOREPTW vpfnSRSetRestorePointW = NULL;
static HMODULE vhSrClientDll = NULL;


static HRESULT InitializeComSecurity();


DAPI_(HRESULT) SrpInitialize(
    __in BOOL fInitializeComSecurity
    )
{
    HRESULT hr = S_OK;

    hr = LoadSystemLibrary(L"srclient.dll", &vhSrClientDll);
    if (FAILED(hr))
    {
        ExitFunction1(hr = E_NOTIMPL);
    }

    vpfnSRSetRestorePointW = reinterpret_cast<PFN_SETRESTOREPTW>(::GetProcAddress(vhSrClientDll, "SRSetRestorePointW"));
    SrpExitOnNullWithLastError(vpfnSRSetRestorePointW, hr, "Failed to find set restore point proc address.");

    // If allowed, initialize COM security to enable NetworkService,
    // LocalService and System to make callbacks to the process
    // calling System Restore. This is required for any process
    // that calls SRSetRestorePoint.
    if (fInitializeComSecurity)
    {
        hr = InitializeComSecurity();
        SrpExitOnFailure(hr, "Failed to initialize security for COM to talk to system restore.");
    }

LExit:
    if (FAILED(hr) && vhSrClientDll)
    {
        SrpUninitialize();
    }

    return hr;
}

DAPI_(void) SrpUninitialize()
{
    if (vhSrClientDll)
    {
        ::FreeLibrary(vhSrClientDll);
        vhSrClientDll = NULL;
        vpfnSRSetRestorePointW = NULL;
    }
}

DAPI_(HRESULT) SrpCreateRestorePoint(
    __in_z LPCWSTR wzApplicationName,
    __in SRP_ACTION action
    )
{
    HRESULT hr = S_OK;
    RESTOREPOINTINFOW restorePoint = { };
    STATEMGRSTATUS status = { };

    if (!vpfnSRSetRestorePointW)
    {
        ExitFunction1(hr = E_NOTIMPL);
    }

    restorePoint.dwEventType = BEGIN_SYSTEM_CHANGE;
    restorePoint.dwRestorePtType = (SRP_ACTION_INSTALL == action) ? APPLICATION_INSTALL : (SRP_ACTION_UNINSTALL == action) ? APPLICATION_UNINSTALL : MODIFY_SETTINGS;
    ::StringCbCopyW(restorePoint.szDescription, sizeof(restorePoint.szDescription), wzApplicationName);

    if (!vpfnSRSetRestorePointW(&restorePoint, &status))
    {
        SrpExitOnWin32Error(status.nStatus, hr, "Failed to create system restore point.");
    }

LExit:
    return hr;
}


// internal functions.

static HRESULT InitializeComSecurity()
{
    HRESULT hr = S_OK;
    DWORD er = ERROR_SUCCESS;
    SECURITY_DESCRIPTOR sd = {0};
    EXPLICIT_ACCESS ea[5] = {0};
    ACL* pAcl = NULL;
    ULONGLONG rgSidBA[(SECURITY_MAX_SID_SIZE+sizeof(ULONGLONG)-1)/sizeof(ULONGLONG)]={0};
    ULONGLONG rgSidLS[(SECURITY_MAX_SID_SIZE+sizeof(ULONGLONG)-1)/sizeof(ULONGLONG)]={0};
    ULONGLONG rgSidNS[(SECURITY_MAX_SID_SIZE+sizeof(ULONGLONG)-1)/sizeof(ULONGLONG)]={0};
    ULONGLONG rgSidPS[(SECURITY_MAX_SID_SIZE+sizeof(ULONGLONG)-1)/sizeof(ULONGLONG)]={0};
    ULONGLONG rgSidSY[(SECURITY_MAX_SID_SIZE+sizeof(ULONGLONG)-1)/sizeof(ULONGLONG)]={0};
    DWORD cbSid = 0;

    // Create the security descriptor explicitly as follows because
    // CoInitializeSecurity() will not accept the relative security descriptors
    // returned by ConvertStringSecurityDescriptorToSecurityDescriptor().
    //
    // The result is a security descriptor that is equivalent to the following 
    // security descriptor definition language (SDDL) string:
    //
    //   O:BAG:BAD:(A;;0x1;;;LS)(A;;0x1;;;NS)(A;;0x1;;;PS)(A;;0x1;;;SY)(A;;0x1;;;BA)
    //
 
    // Initialize the security descriptor.
    if (!::InitializeSecurityDescriptor(&sd, SECURITY_DESCRIPTOR_REVISION))
    {
        SrpExitWithLastError(hr, "Failed to initialize security descriptor for system restore.");
    }

    // Create an administrator group security identifier (SID).
    cbSid = sizeof(rgSidBA);
    if (!::CreateWellKnownSid(WinBuiltinAdministratorsSid, NULL, rgSidBA, &cbSid))
    {
        SrpExitWithLastError(hr, "Failed to create administrator SID for system restore.");
    }

    // Create a local service security identifier (SID).
    cbSid = sizeof(rgSidLS);
    if (!::CreateWellKnownSid(WinLocalServiceSid, NULL, rgSidLS, &cbSid))
    {
        SrpExitWithLastError(hr, "Failed to create local service SID for system restore.");
    }

    // Create a network service security identifier (SID).
    cbSid = sizeof(rgSidNS);
    if (!::CreateWellKnownSid(WinNetworkServiceSid, NULL, rgSidNS, &cbSid))
    {
        SrpExitWithLastError(hr, "Failed to create network service SID for system restore.");
    }

    // Create a personal account security identifier (SID).
    cbSid = sizeof(rgSidPS);
    if (!::CreateWellKnownSid(WinSelfSid, NULL, rgSidPS, &cbSid))
    {
        SrpExitWithLastError(hr, "Failed to create self SID for system restore.");
    }

    // Create a local service security identifier (SID).
    cbSid = sizeof(rgSidSY);
    if (!::CreateWellKnownSid(WinLocalSystemSid, NULL, rgSidSY, &cbSid))
    {
        SrpExitWithLastError(hr, "Failed to create local system SID for system restore.");
    }

    // Setup the access control entries (ACE) for COM. COM_RIGHTS_EXECUTE and
    // COM_RIGHTS_EXECUTE_LOCAL are the minimum access rights required.
    ea[0].grfAccessPermissions = COM_RIGHTS_EXECUTE | COM_RIGHTS_EXECUTE_LOCAL;
    ea[0].grfAccessMode = SET_ACCESS;
    ea[0].grfInheritance = NO_INHERITANCE;
    ea[0].Trustee.pMultipleTrustee = NULL;
    ea[0].Trustee.MultipleTrusteeOperation = NO_MULTIPLE_TRUSTEE;
    ea[0].Trustee.TrusteeForm = TRUSTEE_IS_SID;
    ea[0].Trustee.TrusteeType = TRUSTEE_IS_GROUP;
    ea[0].Trustee.ptstrName = (LPTSTR)rgSidBA;

    ea[1].grfAccessPermissions = COM_RIGHTS_EXECUTE | COM_RIGHTS_EXECUTE_LOCAL;
    ea[1].grfAccessMode = SET_ACCESS;
    ea[1].grfInheritance = NO_INHERITANCE;
    ea[1].Trustee.pMultipleTrustee = NULL;
    ea[1].Trustee.MultipleTrusteeOperation = NO_MULTIPLE_TRUSTEE;
    ea[1].Trustee.TrusteeForm = TRUSTEE_IS_SID;
    ea[1].Trustee.TrusteeType = TRUSTEE_IS_GROUP;
    ea[1].Trustee.ptstrName = (LPTSTR)rgSidLS;

    ea[2].grfAccessPermissions = COM_RIGHTS_EXECUTE | COM_RIGHTS_EXECUTE_LOCAL;
    ea[2].grfAccessMode = SET_ACCESS;
    ea[2].grfInheritance = NO_INHERITANCE;
    ea[2].Trustee.pMultipleTrustee = NULL;
    ea[2].Trustee.MultipleTrusteeOperation = NO_MULTIPLE_TRUSTEE;
    ea[2].Trustee.TrusteeForm = TRUSTEE_IS_SID;
    ea[2].Trustee.TrusteeType = TRUSTEE_IS_GROUP;
    ea[2].Trustee.ptstrName = (LPTSTR)rgSidNS;

    ea[3].grfAccessPermissions = COM_RIGHTS_EXECUTE | COM_RIGHTS_EXECUTE_LOCAL;
    ea[3].grfAccessMode = SET_ACCESS;
    ea[3].grfInheritance = NO_INHERITANCE;
    ea[3].Trustee.pMultipleTrustee = NULL;
    ea[3].Trustee.MultipleTrusteeOperation = NO_MULTIPLE_TRUSTEE;
    ea[3].Trustee.TrusteeForm = TRUSTEE_IS_SID;
    ea[3].Trustee.TrusteeType = TRUSTEE_IS_GROUP;
    ea[3].Trustee.ptstrName = (LPTSTR)rgSidPS;

    ea[4].grfAccessPermissions = COM_RIGHTS_EXECUTE | COM_RIGHTS_EXECUTE_LOCAL;
    ea[4].grfAccessMode = SET_ACCESS;
    ea[4].grfInheritance = NO_INHERITANCE;
    ea[4].Trustee.pMultipleTrustee = NULL;
    ea[4].Trustee.MultipleTrusteeOperation = NO_MULTIPLE_TRUSTEE;
    ea[4].Trustee.TrusteeForm = TRUSTEE_IS_SID;
    ea[4].Trustee.TrusteeType = TRUSTEE_IS_GROUP;
    ea[4].Trustee.ptstrName = (LPTSTR)rgSidSY;

    // Create an access control list (ACL) using this ACE list.
    er = ::SetEntriesInAcl(countof(ea), ea, NULL, &pAcl);
    SrpExitOnWin32Error(er, hr, "Failed to create ACL for system restore.");

    // Set the security descriptor owner to Administrators.
    if (!::SetSecurityDescriptorOwner(&sd, rgSidBA, FALSE))
    {
        SrpExitWithLastError(hr, "Failed to set administrators owner for system restore.");
    }

    // Set the security descriptor group to Administrators.
    if (!::SetSecurityDescriptorGroup(&sd, rgSidBA, FALSE))
    {
        SrpExitWithLastError(hr, "Failed to set administrators group access for system restore.");
    }

    // Set the discretionary access control list (DACL) to the ACL.
    if (!::SetSecurityDescriptorDacl(&sd, TRUE, pAcl, FALSE))
    {
        SrpExitWithLastError(hr, "Failed to set DACL for system restore.");
    }

    // Note that an explicit security descriptor is being passed in.
    hr= ::CoInitializeSecurity(&sd, -1, NULL, NULL, RPC_C_AUTHN_LEVEL_PKT_PRIVACY, RPC_C_IMP_LEVEL_IDENTIFY, NULL, EOAC_DISABLE_AAA | EOAC_NO_CUSTOM_MARSHAL, NULL);
    SrpExitOnFailure(hr, "Failed to initialize COM security for system restore.");

LExit:
    if (pAcl)
    {
        ::LocalFree(pAcl);
    }

    return hr;
}