#include <Windows.h>
#include "detours.h"
#include <string>
#include <wchar.h>
#include <filesystem>
#include <Psapi.h>

void ArgvQuote(const std::wstring&, std::wstring&, bool);
void ShowError(DWORD);
void ShowLastError();

int __stdcall wWinMain(
	_In_ HINSTANCE hInstance,
	_In_opt_ HINSTANCE hPrevInstance,
	_In_ LPWSTR    lpCmdLine,
	_In_ int       nShowCmd
)
{
	int numArgs = 0;
    
	LPWSTR* args = CommandLineToArgvW(lpCmdLine, &numArgs);

    if (numArgs < 2)
    {
        MessageBox(NULL, L"Provide at least a target and a AUMI", L"ERROR", MB_OK);
        return -1;
    }

    LPWSTR cmd = args[0];
    LPWSTR aumi = args[1];

    numArgs -= 2;
    if (numArgs == 0)
    {
        args = NULL;
    }
    else
    {
        args += 2;
    }

    std::wstring cmdLine;
    for (int i = 0; i < numArgs; i++)
    {
        if (cmdLine.length() > 0) cmdLine.push_back(L' ');
        ArgvQuote(std::wstring(args[i]), cmdLine, false);
    }

    wchar_t* cmdLineC = new wchar_t[cmdLine.length() + 1];
    wcscpy_s(cmdLineC, cmdLine.length() + 1, cmdLine.c_str());

    std::filesystem::path injectDll = std::filesystem::canonical(std::filesystem::path("InjectAUMI_64.dll"));

    std::string injectDllStr = injectDll.string();
    //MessageBoxA(NULL, injectDllStr.c_str(), "Path", MB_OK);

    STARTUPINFO startupInfo;
    PROCESS_INFORMATION processInfo;

    ZeroMemory(&startupInfo, sizeof(startupInfo));
    startupInfo.cb = sizeof(startupInfo);
    ZeroMemory(&processInfo, sizeof(processInfo));

    //if (!DetourCreateProcessWithDllEx(cmd, cmdLineC, NULL, NULL, false, 0, NULL, NULL, &startupInfo, &processInfo, injectDllStr.c_str(), NULL))
    if (!CreateProcess(cmd, cmdLineC, NULL, NULL, false, CREATE_SUSPENDED, NULL, NULL, &startupInfo, &processInfo/*, injectDllStr.c_str(), NULL*/))
    {
        ShowLastError();
        //MessageBox(NULL, L"Failed to DetourCreateProcessWithDllEx", L"Error", MB_OK);
        return -1;
    }

    HANDLE process = processInfo.hProcess;

    /*
    HMODULE injected = LoadLibraryA(injectDllStr.c_str());
    if (injected == NULL)
    {
        ShowLastError();
        return -1;
    }

    FARPROC targetAddr = GetProcAddress(injected, "TestExport");
    if (targetAddr == NULL)
    {
        ShowLastError();
        return -1;
    }

    INT_PTR offset = (INT_PTR)targetAddr - (INT_PTR)injected;
    */

    Sleep(10000);

    //SuspendThread(processInfo.hThread);

    const char* msg = "TestInjected";

    HMODULE* remoteModuleList = new HMODULE[1000];
    DWORD needed = 0;
    if (!EnumProcessModules(process, remoteModuleList, 1000, &needed))
    {
        ShowLastError();
        return -1;
    }
    int numModules = needed / sizeof(HMODULE);

    bool found = false;

    wchar_t* name = new wchar_t[1024];

    

    //for (int tries = 5; tries >= 0; tries--)
    //{
        for (int i = 0; i < numModules; i++)
        {
            GetModuleFileNameEx(process, remoteModuleList[i], name, 1023);
            std::wstring modName(name);

            if (modName.find(L"InjectAUMI") != std::wstring::npos)
            {
                //MessageBox(NULL, L"Found", L"Found", MB_OK);

                const char* msg = "Message from here";
                void* remoteMem = VirtualAllocEx(process, NULL, strlen(msg) + 1, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
                if (remoteMem == NULL)
                {
                    ShowLastError();
                    return -1;
                }
                else
                {
                    MessageBox(NULL, L"Got mem", L"Info", MB_OK);
                }

                if (!WriteProcessMemory(process, remoteMem, msg, strlen(msg) + 1, NULL))
                {
                    ShowLastError();
                    return -1;
                }

                //ResumeThread(processInfo.hThread);

                /*
                HANDLE thread = CreateRemoteThread(process, NULL, 0, (LPTHREAD_START_ROUTINE)((INT_PTR)remoteModuleList[i] + offset), remoteMem, 0, NULL);
                if (thread == NULL)
                {
                    ShowLastError();
                    return -1;
                }

                WaitForSingleObject(thread, INFINITE);

                DWORD ret = 0;
                GetExitCodeThread(thread, &ret);

                if (ret != 0)
                {
                    char err[128];
                    sprintf_s(err, "Error from injected function: %d", ret);
                    MessageBoxA(NULL, err, "Error", MB_OK);
                    return -1;
                }
                */

                found = true;
                break;
            }
        }
    //    if (found) break;
    //    Sleep(1000);
    //}

    //ResumeThread(processInfo.hThread);

    if (!found)
    {
        ResumeThread(processInfo.hThread);
        MessageBox(NULL, L"Couldn't find injected dll", L"Error", MB_OK);
        return -1;
    }
    else
    {
        MessageBox(NULL, L"Was Found", L"Info", MB_OK);
    }

    MessageBox(NULL, L"Done", L"Done", MB_OK);

	return 0;
}



void
ArgvQuote(
    const std::wstring& Argument,
    std::wstring& CommandLine,
    bool Force
)

/*++

Routine Description:

    This routine appends the given argument to a command line such
    that CommandLineToArgvW will return the argument string unchanged.
    Arguments in a command line should be separated by spaces; this
    function does not add these spaces.

Arguments:

    Argument - Supplies the argument to encode.

    CommandLine - Supplies the command line to which we append the encoded argument string.

    Force - Supplies an indication of whether we should quote
            the argument even if it does not contain any characters that would
            ordinarily require quoting.

Return Value:

    None.

Environment:

    Arbitrary.

--*/

{
    //
    // Unless we're told otherwise, don't quote unless we actually
    // need to do so --- hopefully avoid problems if programs won't
    // parse quotes properly
    //

    if (Force == false &&
        Argument.empty() == false &&
        Argument.find_first_of(L" \t\n\v\"") == Argument.npos)
    {
        CommandLine.append(Argument);
    }
    else {
        CommandLine.push_back(L'"');

        for (auto It = Argument.begin(); ; ++It) {
            unsigned NumberBackslashes = 0;

            while (It != Argument.end() && *It == L'\\') {
                ++It;
                ++NumberBackslashes;
            }

            if (It == Argument.end()) {

                //
                // Escape all backslashes, but let the terminating
                // double quotation mark we add below be interpreted
                // as a metacharacter.
                //

                CommandLine.append((size_t)NumberBackslashes * 2, L'\\');
                break;
            }
            else if (*It == L'"') {

                //
                // Escape all backslashes and the following
                // double quotation mark.
                //

                CommandLine.append((size_t)NumberBackslashes * 2 + 1, L'\\');
                CommandLine.push_back(*It);
            }
            else {

                //
                // Backslashes aren't special here.
                //

                CommandLine.append(NumberBackslashes, L'\\');
                CommandLine.push_back(*It);
            }
        }

        CommandLine.push_back(L'"');
    }
}

void ShowLastError()
{
    ShowError(GetLastError());
}

void ShowError(DWORD error)
{
    LPTSTR err;
    FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM, NULL, error, 0, (LPTSTR)&err, 0, NULL);
    MessageBoxW(NULL, err, L"Error", MB_OK);
}