////////////////////////////////////////////////////////////////////
//
// File: 	lock.c
// Project: 	lock
//
// Desc:	this sample [program will try to [un]lock
//		a drive, to prevent removal of it's removable 
//		media. It can also display the lock count on
//		the drive. It might not work for all devices!
//
// Revisions: 	Created 24/12/98
//
// Copyright(C) 1998, Tomas Restrepo. All rights reserved
//
///////////////////////////////////////////////////////////////////

#include <windows.h>
#include <stdio.h>
#include <ctype.h>

// DIOC code
#if !defined (VWIN32_DIOC_DOS_IOCTL) 
#define VWIN32_DIOC_DOS_IOCTL      1


// DICO registers, used when calling a int21h call
typedef struct _DIOC_REGISTERS 
{ 
    DWORD reg_EBX;
    DWORD reg_EDX;
    DWORD reg_ECX;
    DWORD reg_EAX;
    DWORD reg_EDI;
    DWORD reg_ESI;
    DWORD reg_Flags;
} DIOC_REGISTERS, *PDIOC_REGISTERS;
#endif 

// Intel x86 processor status flags
#define CARRY_FLAG             0x0001


// All MS-DOS  structuire in IO-CALSS should be aligned on
// 1 byte boundaries.
#pragma pack(1) 
typedef struct _PARAMBLOCK 
{ 
   BYTE bOperation;
   BYTE bNumLocks;
} PARAMBLOCK, *PPARAMBLOCK; 
#pragma pack() 

//////////////////////////////////////////////////////////////
// DisplayLockStatus ()
// 
// Will  display the lock count for the given drive. 
// The bDrive parameter is specified in standard
// DOS drive numbering: 0 - Default drive
//						1 - A
//						2 - B.....
BOOL DisplayLockStatus ( BYTE bDrive )
{
	DIOC_REGISTERS	regs = {0};
	PARAMBLOCK		pblock = {0};
	HANDLE			hDev = NULL;	// handle to VWIN32
	BOOL			fResult;
	DWORD			cb;

	// Open VWIN32 VxD
	hDev = CreateFile ( "\\\\.\\VWIN32", 0, 0, NULL, 0, FILE_FLAG_DELETE_ON_CLOSE, NULL );
	if ( hDev == INVALID_HANDLE_VALUE )
		return FALSE;

	// prepare registers
	pblock.bOperation = 2;   // get lock status
	regs.reg_EAX = 0x440D;
	regs.reg_EBX = bDrive;
	regs.reg_ECX = MAKEWORD(0x48, 0x08);
	regs.reg_EDX = (DWORD)&pblock;

	fResult = DeviceIoControl ( hDev, VWIN32_DIOC_DOS_IOCTL,
								&regs, sizeof(regs), 
								&regs, sizeof(regs),
								&cb, 0 );

	if (fResult)
	{
        // DeviceIoControl succeeded. Now see if the qurey succeeded. It
        // succeeded if the carry flag is not set, or if the carry flag is
        // set but EAX is 0x01 or 0xB0.
		//
        // It failed if the carry flag is set and EAX is not 0x01 or 0xB0.
		//
        // If the carry flag is clear, then unlock succeeded. However, you
        // don't need to set fResult because it is already TRUE when you get
        // in here.

		if (regs.reg_Flags & CARRY_FLAG)
			fResult = (regs.reg_EAX == 0xB0) || (regs.reg_EAX == 0x01);
	}

	// Unlock the drive and release the VxD
	LockLogicalVolume ( hDev, bDrive, FALSE );
	CloseHandle ( hDev );

	if ( fResult )
		// succeeded. Print locks pending
		printf ( "Drive has %d locks pending\n", pblock.bNumLocks );
	

	return fResult;
} 

//////////////////////////////////////////////////////////////
// LockDrive()
// 
// Will try to [un]lock the drive.
// The bDrive parameter is specified in standard
// DOS drive numbering: 0 - Default drive
//						1 - A
//						2 - B.....
// The bLock parameter identifies is locking or unlocking
// is desired:
//			TRUE  - Lock drive
//			FALSE - Unlock drive
BOOL LockDrive ( BYTE bDrive, 
				 BOOL bLock ) 
{ 
	DIOC_REGISTERS	regs = {0};
	PARAMBLOCK		pblock = {0};
	HANDLE			hDev = NULL;	// handle to VWIN32
	BOOL			fResult;
	DWORD			cb;

	// Open VWIN32 VxD
	hDev = CreateFile ( "\\\\.\\VWIN32", 0, 0, NULL, 0, FILE_FLAG_DELETE_ON_CLOSE, NULL );
	if ( hDev == INVALID_HANDLE_VALUE )
		return FALSE;

	// prepare registers
	pblock.bOperation = bLock?0:1;   // lock or unlock?
	regs.reg_EAX = 0x440D;
	regs.reg_EBX = bDrive;
	regs.reg_ECX = MAKEWORD(0x48, 0x08);
	regs.reg_EDX = (DWORD)&pblock;

	fResult = DeviceIoControl ( hDev, VWIN32_DIOC_DOS_IOCTL,
					&regs, sizeof(regs), 
					&regs, sizeof(regs),
					&cb, 0 );

	// See if DeviceIoControl and the lock succeeded
	fResult = fResult && !(regs.reg_Flags & CARRY_FLAG);

	// Unlock the drive and release the VxD
	LockLogicalVolume ( hDev, bDrive, FALSE );
	CloseHandle ( hDev );

	if ( fResult )
		printf ( "Drive has been [un]locked successfully\n" );

	return fResult;

} 


//////////////////////////////////////////////////////////////
// DisplayUsage()
//
// Self-explanatory :)
void DisplayUsage ( void )
{
	puts ( "USAGE: lock <drive> [-l | -u | -s]\n" );
	puts ( "-l: lock drive" );
	puts ( "-u: unlock drive" );
	puts ( "-s: display lock status for drive" );
}

///////////////////////////////////////////////////////////////
// main()
//
// See DisplayUsage() for parameter information.
int main (int argc, char **argv)

{ 
   BYTE		bDrive;
   BOOL		fResult;

   if ( argc != 3 )
   {
	   DisplayUsage ();
	   return 0;
   }
   // convert command line arg 1 from a drive letter to a DOS drive
   // number
   bDrive = (toupper (argv[1][0]) - 'A') + 1;
   
   // check parameters
   if ( argv[2][1] == 'l' )	
	   fResult = LockDrive ( bDrive, TRUE );
   else if ( argv[2][1] == 'u' )
	   fResult = LockDrive ( bDrive, FALSE );
   else if ( argv[2][1] == 's' )
	   fResult = DisplayLockStatus ( bDrive );
   else
   {   DisplayUsage (); return 0; }

   if ( !fResult )
	   printf ( "Failed to [un]lock drive %s\n", argv[1] );
   return 0;

}