// compile with exceptions
// link with user32 and shell32
#ifdef _MSC_VER
#   pragma comment(lib, "user32.lib")
#   pragma comment(lib, "shell32.lib")
#endif
#define _WIN32_WINNT 0x0501
#define OEMRESOURCE
#include <windows.h>
#include <boost/lexical_cast.hpp>

//
// forward decls
//
LRESULT CALLBACK windowProc(HWND, UINT, WPARAM, LPARAM);
void fail(LPTSTR);
void registerClass(HINSTANCE, HICON);
void createWindow(HINSTANCE);
BOOL CALLBACK findWindow(HWND, LPARAM);
HANDLE getTargetProcess(DWORD&);
HICON getTargetIcon(void);

int WINAPI WinMain(HINSTANCE instance, HINSTANCE, LPSTR, int) {
    // get the icon of given running process
    HICON icon = getTargetIcon();
    
    // create our own window with that icon
    registerClass(instance, icon);
    createWindow(instance);
    
    MSG msg;
    while (GetMessage(&msg, NULL, 0, 0)) {
        TranslateMessage(&msg);
        DispatchMessage(&msg);
    }
    
    return 0;
}

//
// our window
//
LRESULT CALLBACK windowProc(HWND wnd, UINT msg, WPARAM wp, LPARAM lp) {
    if (msg == WM_DESTROY) {
        PostQuitMessage(0);
        return 0;
    }
    
    return DefWindowProc(wnd, msg, wp, lp);
}

void registerClass(HINSTANCE instance, HICON icon) {
    WNDCLASSEX wndclass;
    ZeroMemory(&wndclass, sizeof(WNDCLASSEX));
    
    wndclass.cbSize        = sizeof(WNDCLASSEX);
    wndclass.lpfnWndProc   = windowProc;
    wndclass.hInstance     = instance;
    wndclass.hIcon         = icon;
    wndclass.hbrBackground = reinterpret_cast<HBRUSH>(COLOR_BACKGROUND);
    wndclass.hCursor       = static_cast<HCURSOR>(LoadImage(
        NULL, MAKEINTRESOURCE(OCR_NORMAL), IMAGE_CURSOR, 0, 0, LR_DEFAULTSIZE | LR_SHARED
    ));
    wndclass.lpszClassName = TEXT("icon_stealer");
    
    ATOM wndclassAtom = RegisterClassEx(&wndclass);
    if (!wndclassAtom) fail(TEXT("registerclass"));
}

void createWindow(HINSTANCE instance) {
    HWND wnd = CreateWindowEx(
        WS_EX_APPWINDOW, TEXT("icon_stealer"), TEXT("Icon Stealer"), WS_CAPTION | WS_VISIBLE | WS_SYSMENU,
        CW_USEDEFAULT, CW_USEDEFAULT, 150, 40, NULL, NULL, instance, NULL
    );
    
    if (!wnd) fail(TEXT("createwindow"));
    UpdateWindow(wnd);
}

//
// their window
//
HANDLE getTargetProcess(DWORD &pid) {
    int argc;
    LPWSTR *argv;
    
    argv = CommandLineToArgvW(GetCommandLineW(), &argc);
    if (!argv || argc < 2) fail(TEXT("argv"));
    
    pid = boost::lexical_cast<DWORD>(argv[1]);
    
    HANDLE process = OpenProcess(PROCESS_ALL_ACCESS, FALSE, pid);
    if (!process) fail(TEXT("openprocess"));
    
    return process;
}

namespace globals {
    HWND  result;
    DWORD pid;
}

BOOL CALLBACK findWindow(HWND wnd, LPARAM) {
    DWORD foundPid;
    GetWindowThreadProcessId(wnd, &foundPid);
    
    if (foundPid == globals::pid) {
        globals::result = wnd;
        return FALSE;
    }
    
    return TRUE;
}

HICON getTargetIcon(void) {
    HANDLE process = getTargetProcess(globals::pid);
    EnumWindows(findWindow, 0);
    
    return reinterpret_cast<HICON>(GetClassLong(globals::result, GCL_HICON));
}

//
// util
//
void fail(LPTSTR msg) {
    MessageBox(NULL, msg, TEXT("Fail."), MB_OK | MB_ICONERROR);
    ExitProcess(1);
}