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

// For new visual style for MessageBox
#pragma comment(linker,"\"/manifestdependency:type='win32' name='Microsoft.Windows.Common-Controls' version='6.0.0.0' processorArchitecture='*'    publicKeyToken='6595b64144ccf1df' language='*'\"")

static const char* VERSION = "1.1.2";

static const wchar_t* DLL_FILE_PATH = L"SpotifyKeyDumper.dll";
static const wchar_t* PROC_NAME = L"Spotify.exe";

DWORD GetProcId(const wchar_t* procName)
{
    DWORD procId = 0;
    HANDLE hSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

    if (hSnap != INVALID_HANDLE_VALUE)
    {
        PROCESSENTRY32 procEntry = {};
        procEntry.dwSize = sizeof(procEntry);

        if (Process32First(hSnap, &procEntry))
        {
            do
            {
                if (!wcscmp(procEntry.szExeFile, procName))
                {
                    procId = procEntry.th32ProcessID;
                    break;
                }
            }
            while (Process32Next(hSnap, &procEntry));
        }
    }

    CloseHandle(hSnap);
    return procId;
}

void StartSuspendedInjection()
{
    STARTUPINFO startupInfo;
    PROCESS_INFORMATION procInfo;
    void* tmpPage;
    HANDLE injectThread;
    size_t dllPathLen = (wcslen(DLL_FILE_PATH) + 1) * sizeof(wchar_t);

    // Check if Spotify is already open
    DWORD procId = GetProcId(PROC_NAME);
    if (procId)
    {
        std::wstring procNameW = std::wstring(PROC_NAME);
        MessageBox(NULL, std::wstring(procNameW + std::wstring(L" has already been detected!\n\n"
            "Please launch this before ") + procNameW).c_str(), NULL, MB_OK | MB_ICONERROR);
        return;
    }

    std::wcout << "Starting " << PROC_NAME << "..." << std::endl;

    ZeroMemory(&startupInfo, sizeof(startupInfo));
    startupInfo.cb = sizeof(STARTUPINFOA);
    ZeroMemory(&procInfo, sizeof(procInfo));

    if (!CreateProcessW(PROC_NAME, NULL, NULL, NULL, false, CREATE_SUSPENDED, NULL, NULL, &startupInfo, &procInfo))
    {
        std::wcout << "Error: Could not start " << PROC_NAME << " (could not create process): " << GetLastError()
            << std::endl;
        return;
    }

    std::cout << "Injecting DLL..." << std::endl;

    if (GetFileAttributesW(PROC_NAME) == INVALID_FILE_ATTRIBUTES)
    {
        std::wcout << "Error: DLL injection failed (could not find " << PROC_NAME << ")" << std::endl;
        return;
    }

	if (GetFileAttributesW(DLL_FILE_PATH) == INVALID_FILE_ATTRIBUTES)
	{
		std::wcout << "Error: DLL injection failed (could not find " << DLL_FILE_PATH << ")" << std::endl;
		return;
	}

    tmpPage = VirtualAllocEx(procInfo.hProcess, NULL, dllPathLen, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
    if (tmpPage == NULL)
    {
        std::wcout << "Error: DLL injection failed (could not allocate memory): " << GetLastError() << std::endl;
        return;
    }

    if (!WriteProcessMemory(procInfo.hProcess, tmpPage, (PVOID) DLL_FILE_PATH, dllPathLen, NULL))
    {
        std::wcout << "Error: DLL injection failed (could not write memory): " << GetLastError() << std::endl;
        return;
    }

    injectThread = CreateRemoteThread(procInfo.hProcess, NULL, 0, (LPTHREAD_START_ROUTINE)LoadLibrary, tmpPage, 0,
        NULL);
    if (injectThread == NULL)
    {
        std::wcout << "Error: DLL injection failed (could not load library): " << GetLastError() << std::endl;
        return;
    }

    if (WaitForSingleObject(injectThread, UINT_MAX) == WAIT_FAILED)
    {
        std::wcout << "Error: DLL injection failed (could not wait for thread to return): " << GetLastError()
            << std::endl;
        return;
    }

    if (ResumeThread(procInfo.hThread) == -1)
    {
        std::wcout << "Error: DLL injection failed (could not resume thread): " << GetLastError() << std::endl;
        return;
    }

    VirtualFreeEx(procInfo.hProcess, tmpPage, 0, MEM_RELEASE);
    CloseHandle(injectThread);
    CloseHandle(procInfo.hProcess);
    CloseHandle(procInfo.hThread);

    std::wcout << "Finished injecting." << std::endl;
}

int __cdecl main()
{
    std::wcout << "SpotifyKeyDumperInjector v" << VERSION << std::endl << std::endl;
    std::wcout << "Attempting to inject \"" << DLL_FILE_PATH << "\" into \"" << PROC_NAME << "\"..." << std::endl;

    StartSuspendedInjection();

    return 0;
}