[ WIN ] APP Init DLL

class cAppInitDllHelper
{

public:
	cAppInitDllHelper();
	virtual ~cAppInitDllHelper();

	bool AddDllPathA( const std::string& _ssDllPath );
	bool RemoveDllPathA( const std::string& _ssDllPath );
	bool AddDllPathW( const std::wstring& _ssDllPath );
	bool RemoveDllPathW( const std::wstring& _ssDllPath );

private:
	bool IsDllPathPresentA( const std::string& _ssCcurrentValue, const std::string& _ssDllPath );
	bool IsDllPathPresentW( const std::wstring& _ssCcurrentValue, const std::wstring& _ssDllPath );

};


#include "pch.h"
#include "cAppInitDllHelper.h"

#include <iostream>
#include <string>
#include <Windows.h>

namespace {
    const std::string ssRegPathA  = "SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion\\Windows";
    const std::string ssRegValueA = "AppInit_DLLs";
    const std::wstring ssRegPathW = L"SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion\\Windows";
    const std::wstring ssRegValueW= L"AppInit_DLLs";
}

cAppInitDllHelper::cAppInitDllHelper()
{

}

cAppInitDllHelper::~cAppInitDllHelper()
{
    
}


// 중복 검사 함수
bool cAppInitDllHelper::IsDllPathPresentA( const std::string& _ssCcurrentValue, const std::string& _ssDllPath ) 
{
    size_t pos = _ssCcurrentValue.find( _ssDllPath );
    if( pos == std::string::npos ) 
    {
        return false;  // 경로가 없음
    }
    
    // 경로 확인: 정확한 일치 여부 확인
    size_t endPos = pos + _ssDllPath.length();
    return (pos == 0 || _ssCcurrentValue[pos - 1] == ';') &&
        (endPos == _ssCcurrentValue.length() || _ssCcurrentValue[endPos] == ';');
}

// DLL 경로 추가 함수
bool cAppInitDllHelper::AddDllPathA( const std::string& _ssDllPath )
{
    HKEY hKey;
    if( ERROR_SUCCESS !=  RegOpenKeyExA( HKEY_LOCAL_MACHINE, ssRegPathA.c_str(), 0, KEY_ALL_ACCESS, &hKey) )
    {
        std::cerr << "Failed to open registry key." << std::endl;
        return false;
    }

    DWORD bufferSize = 8192;
    char buffer[8192];
    if( ERROR_SUCCESS != RegQueryValueExA( hKey, ssRegValueA.c_str(), NULL, NULL, (LPBYTE)buffer, &bufferSize)  )
    {
        std::cerr << "Failed to query current value." << std::endl;
        RegCloseKey( hKey );
        return false;
    }

    std::string currentValue = buffer;
    if( !IsDllPathPresentA( currentValue, _ssDllPath ) )  // 중복 검사
    {
        if( !currentValue.empty() && currentValue.back() != ';' ) 
        {
            currentValue += ";";
        }
        currentValue += _ssDllPath;

        if( ERROR_SUCCESS != RegSetValueExA( hKey, ssRegValueA.c_str(), 0, REG_SZ, (const BYTE*)currentValue.c_str(), currentValue.size() + 1 ) )
        {
            std::cerr << "Failed to set new value." << std::endl;
            RegCloseKey( hKey );
            return false;
        }
    }

    RegCloseKey( hKey );
    return true;
}

// DLL 경로 삭제 함수
bool cAppInitDllHelper::RemoveDllPathA( const std::string& _ssDllPath )
{
    HKEY hKey;
    if( ERROR_SUCCESS != RegOpenKeyExA( HKEY_LOCAL_MACHINE, ssRegPathA.c_str(), 0, KEY_ALL_ACCESS, &hKey) )
    {
        std::cerr << "Failed to open registry key." << std::endl;
        return false;
    }

    DWORD bufferSize = 8192;
    char buffer[8192] = { 0, };

    if( ERROR_SUCCESS != RegQueryValueExA( hKey, ssRegValueA.c_str(), NULL, NULL, (LPBYTE)buffer, &bufferSize ) )
    {
        std::cerr << "Failed to query current value." << std::endl;
        RegCloseKey( hKey );
        return false;
    }

    std::string currentValue = buffer;
    if( true == IsDllPathPresentA( currentValue, _ssDllPath ) )   // 중복 검사
    {
        size_t pos = currentValue.find( _ssDllPath );
        size_t endPos = pos + _ssDllPath.length();
        if( endPos < currentValue.length() && currentValue[endPos] == ';' ) endPos++;  // 세미콜론 포함 삭제
        currentValue.erase( pos, endPos - pos );

        if( ERROR_SUCCESS != RegSetValueExA( hKey, ssRegValueA.c_str(), 0, REG_SZ, (const BYTE*)currentValue.c_str(), currentValue.size() + 1 ) )
        {
            std::cerr << "Failed to set new value." << std::endl;
            RegCloseKey( hKey );
            return false;
        }
    }

    RegCloseKey( hKey );
    return true;
}


// 유니코드 버전의 문자열 처리를 위한 wstring 사용
bool cAppInitDllHelper::IsDllPathPresentW( const std::wstring& currentValue, const std::wstring& dllPath ) 
{
    size_t pos = currentValue.find( dllPath );
    if( pos == std::wstring::npos ) 
    {
        return false;  // 경로가 없음
    }
    // 경로 확인: 정확한 일치 여부 확인
    size_t endPos = pos + dllPath.length();
    return (pos == 0 || currentValue[pos - 1] == L';') &&
        (endPos == currentValue.length() || currentValue[endPos] == L';');
}

bool cAppInitDllHelper::AddDllPathW( const std::wstring& _ssDllPath ) 
{
    HKEY hKey;
    if( ERROR_SUCCESS !=  RegOpenKeyExW( HKEY_LOCAL_MACHINE, L"SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion\\Windows", 0, KEY_ALL_ACCESS, &hKey )   ) {
        std::wcerr << L"Failed to open registry key." << std::endl;
        return false;
    }

    DWORD bufferSize = 8192;
    wchar_t buffer[8192];
    if( ERROR_SUCCESS != RegQueryValueExW( hKey, L"AppInit_DLLs", NULL, NULL, (LPBYTE)buffer, &bufferSize ) )
    {
        std::wcerr << L"Failed to query current value." << std::endl;
        RegCloseKey( hKey );
        return false;
    }

    std::wstring currentValue = buffer;
    if( !IsDllPathPresentW( currentValue, _ssDllPath ) ) 
    {
        if( !currentValue.empty() && currentValue.back() != L';' ) 
        {
            currentValue += L";";
        }
        currentValue += _ssDllPath;

        if( ERROR_SUCCESS != RegSetValueExW( hKey, L"AppInit_DLLs", 0, REG_SZ, (const BYTE*)currentValue.c_str(), (currentValue.size() + 1) * sizeof( wchar_t ) ))
        {
            std::wcerr << L"Failed to set new value." << std::endl;
            RegCloseKey( hKey );
            return false;
        }
    }

    RegCloseKey( hKey );
    return true;
}

bool cAppInitDllHelper::RemoveDllPathW( const std::wstring& _ssDllPath )
{
    HKEY hKey;
    if( ERROR_SUCCESS != RegOpenKeyExW( HKEY_LOCAL_MACHINE, L"SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion\\Windows", 0, KEY_ALL_ACCESS, &hKey ) )
    {
        std::wcerr << L"Failed to open registry key." << std::endl;
        return false;
    }

    DWORD bufferSize = 8192;
    wchar_t buffer[8192] = {0, };
    if( ERROR_SUCCESS != RegQueryValueExW( hKey, L"AppInit_DLLs", NULL, NULL, (LPBYTE)buffer, &bufferSize )  ) 
    {
        std::wcerr << L"Failed to query current value." << std::endl;
        RegCloseKey( hKey );
        return false;
    }

    std::wstring currentValue = buffer;
    
    if( IsDllPathPresentW( currentValue, _ssDllPath ) )
    {
        size_t pos = currentValue.find( _ssDllPath );
        size_t endPos = pos + _ssDllPath.length();
        if( endPos < currentValue.length() && currentValue[endPos] == L';' ) endPos++;
        currentValue.erase( pos, endPos - pos );

        if( ERROR_SUCCESS != RegSetValueExW( hKey, L"AppInit_DLLs", 0, REG_SZ, (const BYTE*)currentValue.c_str(), (currentValue.size() + 1) * sizeof( wchar_t ) )  )
        {
            std::wcerr << L"Failed to set new value." << std::endl;
            RegCloseKey( hKey );
            return false;
        }
    }

    RegCloseKey( hKey );
    return true;
}