//====================================================================
//
// Defrag.c
//
// Copyright (C) 1997 Mark Russinovich
//
// This program demonstrates the use of NT 4.0 FAT and NTFS cluster
// movement File System Control functions.
//
//====================================================================
#include <windows.h>
#include <stdio.h>
#include <conio.h>
#include "defrag.h"

//--------------------------------------------------------------------
//                         D E F I N E S
//--------------------------------------------------------------------

//
// Interval at which output is paused (in lines)
//
#define PAUSEINTERVAL	24

//
// Size of the buffer we read file mapping information into.
// The buffer is big enough to hold the 16 bytes that 
// come back at the head of the buffer (the number of entries 
// and the starting virtual cluster), as well as 512 pairs
// of [virtual cluster, logical cluster] pairs.
//
#define	FILEMAPSIZE		(512+2)

//
// Size of the bitmap buffer we pass in. Its large enough to
// hold information for the 16-byte header that's returned
// plus the indicated number of bytes, each of which has 8 bits 
// (imagine that!)
//
#define BITMAPBYTES		4096
#define BITMAPSIZE		(BITMAPBYTES+2*sizeof(ULONGLONG))

//
// Invalid longlong number
//
#define LLINVALID		((ULONGLONG) -1)

//--------------------------------------------------------------------
//                        G L O B A L S
//--------------------------------------------------------------------

//
// Handle for the raw volume that was opened
//
HANDLE		VolumeHandle;	

//
// Buffer to read file mapping information into
//
ULONGLONG	FileMap[ FILEMAPSIZE ];

//
// Buffer thats passed to bitmap function
//
BYTE		BitMap[ BITMAPSIZE ];

//
// Bit shifting array for efficient processing of the bitmap
//
BYTE		BitShift[] = { 1, 2, 4, 8, 16, 32, 64, 128 };


//--------------------------------------------------------------------
//                      F U N C T I O N S
//--------------------------------------------------------------------

//--------------------------------------------------------------------
//
// PrintNtError
//
// Translates an NTDLL error code into its text equivalent. This
// only deals with ones commonly returned by defragmenting FS Control
// commands.
//--------------------------------------------------------------------
void PrintNtError( NTSTATUS Status )
{
	switch( Status ) {
	case STATUS_SUCCESS:
		printf("STATUS_SUCCESS\n\n");
		break;
	case STATUS_INVALID_PARAMETER:
		printf("STATUS_INVALID_PARAMETER\n\n");
		break;
	case STATUS_BUFFER_TOO_SMALL:
		printf("STATUS_BUFFER_TOO_SMALL\n\n");
		break;
	case STATUS_ALREADY_COMMITTED:
		printf("STATUS_ALREADY_COMMITTED\n\n");
		break;
	case STATUS_INVALID_DEVICE_REQUEST:
		printf("STATUS_INVALID_DEVICE_REQUEST\n\n");
		break;
	default:
		printf("0x%08x\n\n", Status );
		break;
	}		  
}


//--------------------------------------------------------------------
//
// PrintWin32Error
// 
// Translates a Win32 error into a text equivalent
//
//--------------------------------------------------------------------
void PrintWin32Error( DWORD ErrorCode )
{
	LPVOID lpMsgBuf;
 
	FormatMessage( FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM,
					NULL, ErrorCode, 
					MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
					(LPTSTR) &lpMsgBuf, 0, NULL );
	printf("%s\n", lpMsgBuf );
	LocalFree( lpMsgBuf );
}


//--------------------------------------------------------------------
//
// PrintHelp
//
//--------------------------------------------------------------------
void PrintHelp()
{
	printf("\nCommands:\n\n");
	printf("Dump bitmap free clusters:\n");
	printf("     b [offset]\n");
	printf("Enumerate clusters of file:\n");
	printf("     n [filename]\n");
	printf("Move clusters:\n");
	printf("     m [file] [off] [tgt] [numclust]\n");
	printf("Quit:\n");
	printf("     q\n");
	printf("\n");
}

//--------------------------------------------------------------------
//
// PauseOutput
//
// After n lines have printed, stop and wait for the user to continue.
// 'q' causes the function to return false.
//
//--------------------------------------------------------------------
BOOL PauseOutput( DWORD Count ) 
{
	char	key;

	if( !(Count % PAUSEINTERVAL )) {
		printf("more ('q' to quit): "); 
		fflush(stdout);
		key = getch();
		printf("\n");
		if( key == 'q' ) {
			printf("\nEnumeration aborted.\n\n");
			return FALSE;
		}
	}
	return TRUE;
}


//--------------------------------------------------------------------
//
// OpenVolume
//
// Open the volume for defragging, a flag that is new for NT 4.0.
//
//--------------------------------------------------------------------
DWORD OpenVolume( int DriveId ) 
{
	static char			volumeName[] = "\\\\.\\A:";

	//
	// open the volume
	//
	volumeName[4] = DriveId + 'A'; 
	VolumeHandle = CreateFile( volumeName, GENERIC_READ, 
					FILE_SHARE_READ|FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 
					0, 0 );
	if( VolumeHandle == INVALID_HANDLE_VALUE )	{

		return GetLastError();
	}

	return ERROR_SUCCESS;
}


//--------------------------------------------------------------------
//
// DumpBitmap
//
// Start at the offset specified (if any) and dumps all the free
// clusters on the volume to the end of the volume or until
// the user stops the dump with a 'q'.
//
//--------------------------------------------------------------------
void DumpBitmap( char *argument )
{
	DWORD						status;
	PBITMAP_DESCRIPTOR			bitMappings;
	ULONGLONG					cluster;
	ULONGLONG					numFree;
	ULONGLONG					startLcn;
	ULONGLONG					nextLcn;
	ULONGLONG					lastLcn;
	IO_STATUS_BLOCK				ioStatus;
	ULONGLONG					i;
	int							lines;

	//
	// Start scanning at the cluster offset the user specifies
	//
	bitMappings = (PBITMAP_DESCRIPTOR) BitMap;
	cluster = 0;
	nextLcn = 0;
	lines = 0;
	lastLcn = LLINVALID;
	sscanf( argument, " %I64d ", &nextLcn );
	while( !(status = NtFsControlFile( VolumeHandle, NULL, NULL, 0, &ioStatus,
						FSCTL_GET_VOLUME_BITMAP,
						&nextLcn, sizeof( cluster ),
						bitMappings, BITMAPSIZE )) ||
			 status == STATUS_BUFFER_OVERFLOW ||
			 status == STATUS_PENDING ) {

		// 
		// If the operation is pending, wait for it to finish
		//
		if( status == STATUS_PENDING ) {
			
			WaitForSingleObject( VolumeHandle, INFINITE );
			
			//
			// Get the status from the status block
			//
			if( ioStatus.Status != STATUS_SUCCESS && 
				ioStatus.Status != STATUS_BUFFER_OVERFLOW ) {

				printf("\nGet Volume Bitmap: ");
				PrintNtError( ioStatus.Status );
				return;
			}
		}

		//
		// Print the range we're starting at
		//
		if( !lines ) {
		
			printf("\nFree clusters starting at offset: %I64d\n", 
						bitMappings->StartLcn );
		}

		//
		// Scan through the returned bitmap info, looking for empty clusters
		//
		startLcn = bitMappings->StartLcn;
		numFree = 0;
		cluster = LLINVALID;
		for( i = 0; i < min( bitMappings->ClustersToEndOfVol, 8*BITMAPBYTES); i++ ) {

			if( !(bitMappings->Map[ i/8 ] & BitShift[ i % 8 ])) {

				//
				// Cluster is free
				//
				if( cluster == LLINVALID ) {

					cluster = startLcn + i;
					numFree = 1;

				} else {

					numFree++;
				}
			} else {

				//
				// Cluster is not free
				//
				if( cluster != LLINVALID ) {
					
					if( lastLcn == cluster ) {

						lastLcn = LLINVALID;
					} else {

						//
						// See if we should continue
						//
						if( !PauseOutput( ++lines ) ) {
							
							return;
						}

						printf("   LCN: %I64d LEN: %I64d\n", cluster, numFree );
						numFree = 0;
						lastLcn = cluster;
						cluster = LLINVALID;
					}
				} 
			}
		}

		//
		// See if we should continue
		//
		if( !PauseOutput( ++lines ) ) {
			
			return;
		}

		//
		// Print any remaining
		//
		if( cluster != LLINVALID && lastLcn != cluster ) {
			
			printf("   LCN: %I64d LEN: %I64d\n", cluster, numFree );
			numFree = 0;
			cluster = LLINVALID;
		}

		//
		// End of volume?
		//
		if( status != STATUS_BUFFER_OVERFLOW ) {

			printf("End of volume.\n\n");
			return;
		}				

		//
		// Move to the next block
		//
		nextLcn = bitMappings->StartLcn + i;
	}

	//
	// We only get here when there's an error
	//
	printf("\nGet Volume Bitmap: ");
	PrintNtError( status );
}


//--------------------------------------------------------------------
//
// DumpFile
//
// Dumps the clusters belonging to the specified file until the
// end of the file or the user stops the dump.
//
//--------------------------------------------------------------------
void DumpFile( int drive, char *argument )
{
	DWORD						status;
	int							i;
	HANDLE						sourceFile;
	char						fileName[MAX_PATH];
	IO_STATUS_BLOCK				ioStatus;
	ULONGLONG					startVcn;
	PGET_RETRIEVAL_DESCRIPTOR	fileMappings;
	int							lines = 0;

	//
	// Make the name into a real pathname
	//
	if( strlen( argument ) > 1 && argument[0] != '\\' &&
		argument[0] != 'A'+drive &&
		argument[0] != 'a'+drive ) 
		sprintf(fileName, "%C:\\%s", drive+'A', argument );
	else if( strlen( argument ) > 1 && argument[0] == '\\')