#define WIN32_LEAN_AND_MEAN

#ifndef _WIN32_WINNT
#define _WIN32_WINNT 0x0500 // Needed for AttachConsole & Friends
#endif

#include <Windows.h>
#include <Tlhelp32.h>
#include <stdio.h>

#include <fcntl.h>
#include <io.h>


FILE *gpFile = NULL;

void WriteToLog(const char *pPrintMe)
{
    if(gpFile == NULL)
    {
        gpFile = fopen("output.log", "w");
    }

    fprintf(gpFile, pPrintMe);
    fprintf(gpFile, "\n");
}


void ReportError()
{
    DWORD lastError = GetLastError();

    const int kMaxSize = 256;
    LPVOID buffer;

    FormatMessage( FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM
                 , NULL
                 , lastError
                 , MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT)
                 , (LPTSTR)(&buffer)
                 , 0
                 , NULL
                 );

    WriteToLog(static_cast<const char *>(buffer));

    LocalFree(buffer);

}

/**
  * Create a snapshot and test it.
  */
HANDLE CreateSnapshot()
{
    HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

    if(hSnapshot == INVALID_HANDLE_VALUE)
    {
        ReportError();
    }

    return hSnapshot;
}

/**
  * Find a given process by id.
  *
  * @param[in]  hSnapshot     The HANDLE to the snapshot we'll be using.
  * @param[in]  id            The ID of the process to find.
  * @param[out] pProcessEntry The place to store the process entry.
  *
  * @return true if we found the id, false if we passed the end.
  */
bool FindProcessById(HANDLE hSnapshot, const DWORD id, PROCESSENTRY32 *pProcessEntry)
{
    if (Process32First(hSnapshot, pProcessEntry) == FALSE)
    {
        ReportError();
    }

    do 
    {
        if(pProcessEntry->th32ProcessID == id)
        {
            return true;
        }
    } while(Process32Next(hSnapshot, pProcessEntry));
    
    return false;
}

int WINAPI WinMain( HINSTANCE /* handleToCurrentInstance */
                  , HINSTANCE /* handleToPreviousInstance */
                  , LPSTR pCommandLine
                  , int       /* windowDisplayMethod */
                  )
{
    DWORD currentProcessId = GetCurrentProcessId();

    PROCESSENTRY32  processEntry;

    HANDLE hSnapshot = CreateSnapshot();
    if( FindProcessById(hSnapshot, currentProcessId, &processEntry) == false)
    {
        WriteToLog("Couldn't find the current process!");
        return -1;
    }

    if( FindProcessById(hSnapshot, processEntry.th32ParentProcessID, &processEntry) == false)
    {
        WriteToLog("Couldn't find the parent process!");
        return -2;
    }
    
    printf("Parent Process Name = %s", processEntry.szExeFile);

    // OK, phew. Now we finally have our 'parent process' id and info. Lets see if we can attach to that console.
    if(strcmp(processEntry.szExeFile, "cmd.exe") == 0)
    {
        if(AttachConsole(processEntry.th32ProcessID) == FALSE)
        {
            ReportError();
        }

        HANDLE hParentProcess = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, processEntry.th32ProcessID);

        if(hParentProcess == NULL)
        {
            ReportError();
        }

        /**
          * START VOODOO
          * I haven't analyzed this yet, just ripped it from someone. Doesn't work as I'd like.
          */
        intptr_t handleAsInt = (intptr_t)(GetStdHandle(STD_OUTPUT_HANDLE));

        int hCrt = _open_osfhandle(handleAsInt,_O_TEXT);

        FILE *hf = _fdopen( hCrt, "w" );
        *stdout = *hf;
        *stderr = *hf;
        setvbuf( stdout, NULL, _IONBF, 0 );
        setvbuf( stderr, NULL, _IONBF, 0 );

        /**
          * END VOODOO
          * I haven't analyzed this yet, just ripped it from someone. Doesn't work as I'd like.
          */

        printf("Why don't this work? :D");
        WriteToLog("Attached to console successfully!");
    }
    else
    {
        WriteToLog(processEntry.szExeFile);
    }

    printf("Got stuff!");
    CloseHandle(hSnapshot);
    return 0;
}