ModifiedWindowsEnvironment.cpp

#include <windows.h>
#include <cwchar>
#include <string>
#include <map>
#include <iostream>
#include <vector>

// "using" instead of "typedef"s because that's what the cool kids do
using env_string_t = std::wstring;
using env_char_t = wchar_t;
using env_map_t = std::map<env_string_t, env_string_t>;
using env_vector_t = std::vector<env_char_t>;


// env_map_t MapFromWindowsEnvironment()
// * Imports current Environment in a C-string table
// * Parses the strings by splitting on the first "=" per line
// * Creates a map of the variables
// * Returns the map
inline env_map_t MapFromWindowsEnvironment(){
    wchar_t *variable_strings_ptr; 
    wchar_t *environment_strings_ptr;
    std::wstring delimeter(L"=");
    int del_len = delimeter.length();
    env_map_t mapped_environment;

    // Get a pointer to the environment block. 
    environment_strings_ptr = GetEnvironmentStringsW();
    // If the returned pointer is NULL, exit.
    if (environment_strings_ptr == NULL)
    {
        std::cout << "GetEnvironmentStrings failed: " << GetLastError() << std::endl; 
        return mapped_environment;
    }
 
    // Variable strings are separated by NULL byte, and the block is 
    // terminated by a NULL byte. 

    variable_strings_ptr = (wchar_t *) environment_strings_ptr;

    //Since the environment map ends with a null, we can loop until we find it.
    while (*variable_strings_ptr)
    {
        // Create a string from Variable String
        env_string_t current_line(variable_strings_ptr);
        // Find the first "equals" sign.
        auto pos = current_line.find(delimeter);
        // Assuming it's not missing ...
        if(pos!=std::wstring::npos){
            // ... parse the key and value.
            env_string_t key = current_line.substr(0, pos);
            env_string_t value = current_line.substr(pos + del_len);
            // Map the entry.
            mapped_environment[key] = value;
        }
        // Jump to next line in the environment map.
        variable_strings_ptr += std::wcslen(variable_strings_ptr) + 1;
    }
    // We're done with the old environment map buffer.
    FreeEnvironmentStringsW(environment_strings_ptr);

    // Return the map.
    return mapped_environment;
}

// env_vector_t WindowsEnvironmentVectorFromMap(const env_map_t &source_map)
// * Creates a vector buffer for the new environment string table
// * Copies in the mapped variables
// * Returns the vector
inline env_vector_t WindowsEnvironmentVectorFromMap(const env_map_t &source_map)
{
    // Make a new environment map buffer.
    env_vector_t environment_map_buffer;
    // Give it some space.
    environment_map_buffer.reserve(4096);
    
    // And fill'er up.
    for(auto kv: source_map){
        // Create the line
        env_string_t current_line(kv.first); current_line += L"="; current_line += kv.second;
        // Add the line to the buffer.
        std::copy(current_line.begin(), current_line.end(), std::back_inserter(environment_map_buffer));
        // Append a null
        environment_map_buffer.push_back(0);
    }
    // Append one last null because of how Windows does it's environment maps.
    environment_map_buffer.push_back(0);
    
    return environment_map_buffer;
}


// Main program:
// * Imports old environment map
// * Modifies it
// * Exports the modified map to a usable format
// * Spawns a child process with the environment
// * Waits for it to finish
int main(){

    int return_value = 0;

    // Import the environment map
    env_map_t environment_map = MapFromWindowsEnvironment();
    // Setup the changes map
    env_map_t environment_map_changes{{L"WIX", L"Shpadoinkle!"}, {L"WOOT", L"w007"}};

    // Merge in the changes
    // I didn't use "merge" because that only inserts new variables without updating old ones.
    // I didn't merge into the "changes map" because I'm assuming it will be MUCH smaller and
    // it would just be a waste of cycles and memory.
    for(auto& it : environment_map_changes)
    {
        environment_map[it.first] = it.second;
    }

    // Create a Windows-usable Environment Map Buffer
    env_vector_t new_environment_map_strings = WindowsEnvironmentVectorFromMap(environment_map);
    // Grab a pointer to its data
    void* environment_string_table_ptr = (void*)new_environment_map_strings.data();

    // Create the process and wait for it to exit
    {
        DWORD creation_flags = CREATE_UNICODE_ENVIRONMENT;
        STARTUPINFOW startup_info;
        PROCESS_INFORMATION process_info;
        BOOL creation_success;
        SecureZeroMemory(&startup_info, sizeof(STARTUPINFOW));
        startup_info.cb = sizeof(STARTUPINFOW);

        // Attempt to create a child process
        creation_success = CreateProcessW(
            L"C:\\Windows\\System32\\cmd.exe",          // Application name
            NULL,                                       // Command Line
            NULL,                                       // ProcessAttributes
            NULL,                                       // Thread Attributes
            TRUE,                                       // Inherit Handles
            creation_flags,                             // Creation Flags
            environment_string_table_ptr,               // embed new environment data
            NULL,                                       // Current Directory
            &startup_info,                              // Startup Information
            &process_info                               // Process Information
        );

        // Sometimes we fail
        if (! creation_success)
        {
            std::cout<<"CreateProcess failed " << GetLastError() << std::endl;
            return_value = -1;
        }
        // Wait for the child process to finish
        WaitForSingleObject(process_info.hProcess, INFINITE);
    }

    return return_value;
}