/*++


Module Name:

    kloader.c

Abstract:

    Program to load/unload KM driver.

Author:

    Andrey Shedel <andreys@cr.cyco.com>

   You may distribute under the terms of the GNU General Public License
   You can contact author at andreys@cr.cyco.com or http://www.chat.ru/~ashedel

--*/




#define WIN32_LEAN_AND_MEAN
#define _DLL

#include <stdarg.h>
#include <windows.h>
#include <shellapi.h>

#pragma warning(push, 4)

#pragma warning(disable:4204) // nonstandard extension used : non-constant aggregate initializer

//#pragma comment(exestr, "\n\nKLOADER: Kernel-mode driver loader\nCopyright (C) 1999 Andrey Shedel (andreys@tarzan.cr.cyco.com)\n\n")
#pragma comment(lib,"ntdll.lib")
#pragma comment(lib,"kernel32.lib")
#pragma comment(lib,"shell32.lib")
#pragma comment(linker, "-entry:raw_main")
#pragma comment(linker, "-opt:nowin98")
#pragma comment(linker, "-opt:ref")
#pragma comment(linker, "-merge:.rdata=.text")
#pragma comment(linker, "-subsystem:console")



typedef ULONG NTSTATUS;

typedef struct {
	USHORT Length;
	USHORT MaximumLength;
	PWSTR  Buffer;
} UNICODE_STRING, *PUNICODE_STRING;


#define STATUS_SUCCESS 0


#define SE_LOAD_DRIVER_PRIVILEGE          (10L)



NTSYSAPI
NTSTATUS
NTAPI
RtlAppendUnicodeToString (
    PUNICODE_STRING Destination,
    PWSTR Source
    );


NTSYSAPI
NTSTATUS
NTAPI
RtlAdjustPrivilege(
    IN ULONG Privilege,
    IN BOOLEAN NewValue,
    IN BOOLEAN ForThread,
    OUT PBOOLEAN OldValue
   );

NTSYSAPI
NTSTATUS
NTAPI
NtLoadDriver(
    IN PUNICODE_STRING DriverServiceName
   );



NTSYSAPI
NTSTATUS
NTAPI
NtUnloadDriver(
    IN PUNICODE_STRING DriverServiceName
   );




#define SERVICE_PREFIX L"\\Registry\\Machine\\System\\CurrentControlSet\\Services\\"
#define SERVICE_PREFIX_LEN (sizeof(SERVICE_PREFIX) - sizeof(WCHAR))



static const WCHAR szNTDLL_DLL[] = L"ntdll.dll";
static const char szSucceeded[]  = "succeeded.";
static const char szFailed[]     = "failed. ";
static const char szUsage[]      = "KLOADER: Kernel-mode driver loader\nCopyright (C) 1999 Andrey Shedel (andreys@cr.cyco.com)\nUsage: kloader <driver_name> [/r]\n";
static const char szPrivilegeSetError[] = "KLOADER: unable to set privilege\n";
static const WCHAR szServicesPrefix[] = SERVICE_PREFIX;
static const char szMessagePrefix[] = "KLOADER: Attemt ";
static const char szCRLF[]       = "\n";




void
__stdcall
_puts(
    IN const char* msg
   )
{
	ULONG Written;
	WriteFile(GetStdHandle(STD_OUTPUT_HANDLE), msg,
			  strlen(msg), &Written, NULL);
}



__inline
char*
__stdcall
GetErrorString(
    IN NTSTATUS error
   )
{
	char* lpMsgBuf;

	FormatMessageA(
				   FORMAT_MESSAGE_ALLOCATE_BUFFER |
				   FORMAT_MESSAGE_IGNORE_INSERTS |
				   FORMAT_MESSAGE_FROM_SYSTEM |
				   FORMAT_MESSAGE_FROM_HMODULE,
				   GetModuleHandleW(szNTDLL_DLL),
				   error,
				   0,
				   (char*)&lpMsgBuf, 0, NULL);

	return lpMsgBuf;
}





void
__stdcall
pstatus(
    IN NTSTATUS Status
   )
{
	if(STATUS_SUCCESS == Status)
	{
		_puts(szSucceeded);
	}
	else
	{
		char* s = GetErrorString(Status);

		_puts(szFailed);

		if(NULL != s)
		{
			_puts(s);
			LocalFree((HLOCAL)s);
		}
	}
}



__inline
BOOLEAN
__stdcall
EnablePrivilege(
    IN ULONG privilege
   )
{
	BOOLEAN OldValue;
	return (BOOLEAN)(STATUS_SUCCESS == RtlAdjustPrivilege(privilege,
														  TRUE,
														  FALSE,
														  &OldValue
														 ));
}





int
__stdcall
real_wmain(int argc, PCWSTR* argv)
{
	UNICODE_STRING UnicodeString;
	NTSTATUS status;
	WCHAR buffer[1024];


	if(argc < 2)
	{
	wrong_arg:
		_puts(szUsage);
		return 1;
	}


	if(!EnablePrivilege(SE_LOAD_DRIVER_PRIVILEGE))
	{
		_puts(szPrivilegeSetError);
		return 0;
	}


	RtlCopyMemory(&buffer[0], szServicesPrefix, SERVICE_PREFIX_LEN);

	UnicodeString.MaximumLength = sizeof(buffer);
	UnicodeString.Length = SERVICE_PREFIX_LEN;
	UnicodeString.Buffer = &buffer[0];

	RtlAppendUnicodeToString(&UnicodeString, (PWSTR)(argv[1]));


	if(argc < 3)
		status = NtLoadDriver(&UnicodeString);
	else
	{
		if((L'/'  == argv[2][0]) &&
		   (L'r'  == argv[2][1]) &&
		   (L'\0' == argv[2][2]))
		{
			status = NtUnloadDriver(&UnicodeString);
		}
		else
		{
			goto wrong_arg;
		}
	}

	_puts(szMessagePrefix);
	pstatus(status);
	_puts(szCRLF);

	return 0;
}



void __cdecl raw_main(void)
{
	int argc;
	LPWSTR* argv = CommandLineToArgvW(GetCommandLineW(), &argc);

	ExitProcess(real_wmain(argc, argv));
}