C / C++ / MFC

// Image3.cpp : Defines the entry point for the console application.
//

#define UNICODE 1
#define _UNICODE 1
#include "stdafx.h"
#include "windows.h"
#include "stdio.h"
#include "ctype.h"
#include "strsafe.h"
#include "assert.h"

#define ARRLEN(c) (sizeof(c)/sizeof(c[0]))


template<class IntType>
bool AtoI(WCHAR *sz, IntType *out) {
	int i = 0;
	bool sign = false;
	if (sz[i] == '-') {
		++i;
		sign = true;
	}  
	int base = 10;
	if(sz[i] == '0' && tolower(sz[i+1]) == 'x') {
		i += 2;
		base = 16;
	} else if (sz[i] == '0' && sz[i+1] != '\0') {
		++i;
		base = 8;
	}
	IntType res = 0;
	int chars = 0;
	while (sz[i] != '\0') {
		WCHAR c = tolower(sz[i++]);
		if (c >= '0' && c <= '9')
			c = c - '0';
		else if (c >= 'a' && c <= 'f')
			c = c - 'a' + 10;
		else
			return false;
		if (c >= base)
			return false;
		IntType res_sav = res;
		res *= (IntType)base;
		if (res / (IntType)base != res_sav) {
			// Overflow!
			return false;
		}
		res_sav = res;
		res += (IntType)c;
		if (res < res_sav) {
			// Overflow!
			return false;
		}
		++chars;
	}
	if (chars == 0)
		return false;
	*out = sign ? -res : res;
	return true;
}

bool GetDeviceName(WCHAR *szDevName,
                   size_t cchDevName,
                   WCHAR *szInput) {
  unsigned int disk_no = 0;
  if (AtoI<unsigned int>(szInput, &disk_no)) {
    // It's a physical disk
    StringCchPrintfW(szDevName,
                     cchDevName,
                     L"\\\\.\\PHYSICALDRIVE%d", 
                     disk_no);
    return true;
  }

  WCHAR drive_letter = toupper(szInput[0]);
  if (drive_letter >= 'A' && 
      drive_letter <= 'Z' &&
      szInput[1] == ':' &&
      szInput[2] == '\0') {
    // It's a drive letter
    StringCchPrintfW(szDevName,
                     cchDevName,
                     L"\\\\.\\%c:",
                     drive_letter);
    return true;
  }

  WCHAR sz[_MAX_PATH];
  const WCHAR *p = wcsrchr(szInput, '\\');
  if (!p || p[1] != '\0') {
    // Mount point needs to end in backslash
    StringCchCopyW(sz, ARRLEN(sz), szInput);
    StringCchCatW(sz, ARRLEN(sz), L"\\");
    szInput = sz;
  }

  if (!GetVolumeNameForVolumeMountPointW(
      szInput, szDevName, cchDevName))
    return false;

  // This returns a name ending with '\'
  // but CreateFile does not take it, so
  // we need to get rid of it.
  WCHAR *q = wcsrchr(szDevName, '\\');
  if (q && q[1] == '\0')
    *q = '\0';

  return true;
}

HANDLE OpenBlockDevice(const WCHAR *szDevName,
                       bool fWrite) {
  HANDLE h = CreateFile(szDevName,
    fWrite ? GENERIC_WRITE : GENERIC_READ,
    FILE_SHARE_WRITE, NULL, OPEN_EXISTING,
    FILE_FLAG_NO_BUFFERING |
        (fWrite ? FILE_FLAG_WRITE_THROUGH : 0),
    NULL);
  return h;
}

VOLUME_BITMAP_BUFFER *GetVolumeBitmap(
    HANDLE hDev, unsigned int *puiBitmapSize) {
  STARTING_LCN_INPUT_BUFFER sStartLcn;
  sStartLcn.StartingLcn.QuadPart = 0;

  DWORD dwBitmapSize = 0;
  DWORD dwAllocatedSize = 64 * 1024;
  VOLUME_BITMAP_BUFFER *pVolumeBitmap = NULL;
  for ( ; ; ) {
    pVolumeBitmap = (VOLUME_BITMAP_BUFFER *)
        LocalAlloc(LMEM_FIXED, dwAllocatedSize);

    BOOL ret = DeviceIoControl(hDev,
                      FSCTL_GET_VOLUME_BITMAP,
                      &sStartLcn,
                      sizeof(sStartLcn),
                      pVolumeBitmap,
                      dwAllocatedSize,
                      &dwBitmapSize,
                      NULL);
    if (ret) {
      *puiBitmapSize = dwBitmapSize;
      return pVolumeBitmap;
    }

    if (GetLastError() != ERROR_MORE_DATA)
      return NULL;

    dwAllocatedSize *= 2;
  }
}

#define SIG "Compressed Disk Image"

struct BackupHeader {
  char signature[sizeof(SIG)];
  unsigned int uiSectorSize;
  unsigned int uiClusterSize;
  unsigned int uiBitmapSize;
};


//Notice that in the case where pVolumeBitmap is NULL, 
//the function gets reduced to a very simple loop that just reads and writes the entire block device.
// This function assumes the ownership of
// pVolumeBitmap, i. e. it frees it.
// hFrom is the handle from which the data is read.
// hTo is the handle to which the data is written.
// hSeek is always the handle of the device - this
// is the pointer that is being moved to skip empty
// extents.
// For backup: hFrom = hSeek - disk handle
//             hTo - file handle
// For restore: hTo = hSeek - disk handle
//             hFrom - file handle
// If device is not an NTFS volume, pBackupData,
// pVolumeBitmap are both NULL.

DWORD Transfer(HANDLE hFrom,
               HANDLE hTo,
               HANDLE hSeek,
               BackupHeader *pBackupData,
               VOLUME_BITMAP_BUFFER *pVolumeBitmap,
               void *buffer,
               const unsigned int cbBufferSize,
               unsigned __int64 &ui64Bytes) {
  DWORD dwErr = ERROR_SUCCESS;

  // This is used to step through the bitmap
  // of used clusters.
  unsigned char bmpMask = 1;
  unsigned char *bmpIndex = NULL;
  int maxClustersPerRead = 0;
  unsigned __int64 maxClusters = 0;
  if (pVolumeBitmap) {
    bmpIndex = pVolumeBitmap->Buffer;
    maxClusters = pVolumeBitmap->BitmapSize.QuadPart;
    maxClustersPerRead = cbBufferSize /
        pBackupData->uiClusterSize;
  }

  unsigned __int64 cluster = 0;
  bool fTerm = false;

  while (!fTerm) {
    DWORD dwToMove = cbBufferSize;
    if (pVolumeBitmap) {
      // Skip empty. Note that at the end
      // of this loop cluster can be more
      // than maxCluster, because sometimes
      // we step in 8.
      while (cluster < maxClusters) {
        if (!*bmpIndex) {
          cluster += 8;

          ++bmpIndex;
          assert(bmpMask == 1);

          continue;
        }

        if ((*bmpIndex & bmpMask) == 0) {
          ++cluster;

          bmpMask <<= 1;
          if (! bmpMask) {
            bmpMask = 1;
            ++bmpIndex;
          }

          continue;
        }
        break;
      }

      dwToMove = 0;
      int numClusters = 0;
      while (cluster  + numClusters < maxClusters
          && numClusters < maxClustersPerRead
          && (*bmpIndex & bmpMask) != 0) {
        dwToMove += pBackupData->uiClusterSize;
        ++numClusters;

        bmpMask <<= 1;
        if (!bmpMask) {
          bmpMask = 1;
          ++bmpIndex;
        }
      }

      assert(dwToMove <= cbBufferSize);
      if (dwToMove == 0) { // The End?
        // The last sector on the volume
        // contains duplicate of the boot
        // sector, but is not part of the
        // volume map.
        dwToMove = pBackupData->uiSectorSize;

        // Since we might have skipped more
        // clusters than there were because
        // sometimes we count by 8-s, we reset
        // the current cluster count as well:
        cluster = maxClusters;

        // Convenient place to free the bitmap.
        LocalFree(pVolumeBitmap);
        pVolumeBitmap = NULL;
        fTerm = true;
      }

      unsigned __int64 offset = cluster *
        (unsigned __int64)pBackupData->uiClusterSize;

      if (!SetFilePointerEx(hSeek,
                            *(LARGE_INTEGER *)&offset,
                            NULL,
                            FILE_BEGIN)) {
        dwErr = GetLastError();
        wprintf(L"Seek error %d (0x%08x)\n",
            dwErr, dwErr);
      }

      cluster += numClusters;
    }

    DWORD dwRead = 0;
    if (!ReadFile(hFrom, buffer,
        dwToMove, &dwRead, NULL)) {
        dwErr = GetLastError();
        wprintf(L"\nRead error %d (0x%08x)\n",
            dwErr, dwErr);
        break;
    }

    if (dwRead == 0)
      break;

    DWORD dwWrit = 0;
    if (!WriteFile(hTo, buffer,
        dwRead, &dwWrit, NULL) ||
        dwWrit != dwRead) {
      dwErr = GetLastError();
      wprintf(L"\nWrite error %d (0x%08x)\n",
          dwErr, dwErr);
      break;
    }

    ui64Bytes += dwWrit;
    wprintf(L"%I64u\r", ui64Bytes);
  }

  return dwErr;
}
DWORD ReadToFile(HANDLE hDev,
                 HANDLE hFile,
                 void *buffer,
                 const unsigned int cbBufferSize,
                 unsigned __int64 &ui64BytesRead) {
  DWORD dwErr = ERROR_SUCCESS;

  unsigned int uiBitmapSize = 0;
  VOLUME_BITMAP_BUFFER *pVolumeBitmap = NULL;

  NTFS_VOLUME_DATA_BUFFER sVolumeData;
  DWORD dwRead = 0;
  // For NTFS volumes we can store just the
  // clusters that are used. We get the
  // bitmap where 1 means that the cluster
  // is used here.
  if (DeviceIoControl(hDev,
                      FSCTL_GET_NTFS_VOLUME_DATA,
                      NULL,
                      0,
                      &sVolumeData,
                      sizeof(sVolumeData),
                      &dwRead,
                      NULL)) {
    // For NTFS disable volume boundary checks by FS
    DWORD dwBytes = 0;
    DeviceIoControl(hDev,
                    FSCTL_ALLOW_EXTENDED_DASD_IO,
                    NULL,
                    0,
                    NULL,
                    0,
                    &dwBytes,
                    NULL);

    pVolumeBitmap = GetVolumeBitmap(hDev, &uiBitmapSize);
  }

  BackupHeader h;
  DWORD dwWrit = 0;
  if (pVolumeBitmap) {
    memcpy(h.signature, SIG, sizeof(SIG));
    h.uiBitmapSize = uiBitmapSize;
    h.uiSectorSize = sVolumeData.BytesPerSector;
    h.uiClusterSize = sVolumeData.BytesPerCluster;
    if (!WriteFile(hFile,
                   &h,
                   sizeof(h),
                   &dwWrit,
                   NULL)
        || dwWrit != sizeof(h)) {
      dwErr = GetLastError();
      wprintf(L"Write error %d (0x%08x)\n",
          dwErr, dwErr);
      return dwErr;
    }

    if (!WriteFile(hFile,
                   pVolumeBitmap,
                   uiBitmapSize,
                   &dwWrit,
                   NULL)
        || dwWrit != uiBitmapSize) {
      dwErr = GetLastError();
      wprintf(L"Write error %d (0x%08x)\n",
          dwErr, dwErr);
      return dwErr;
    }
  }

  return Transfer(hDev,
                  hFile,
                  hDev,
                  &h,
                  pVolumeBitmap,
                  buffer,
                  cbBufferSize,
                  ui64BytesRead);
}
DWORD WriteFromFile(HANDLE hDev,
                 HANDLE hFile,
                 void *buffer,
                 const unsigned int cbBufferSize,
                 unsigned __int64 &ui64BytesWritten) {
  DWORD dwRead = 0;
  DeviceIoControl(hDev,
                  FSCTL_ALLOW_EXTENDED_DASD_IO,
                  NULL,
                  0,
                  NULL,
                  0,
                  &dwRead,
                  NULL);

  unsigned int uiBitmapSize = 0;
  VOLUME_BITMAP_BUFFER *pVolumeBitmap = NULL;

  BackupHeader h;

  if (ReadFile(hFile, &h, sizeof(h), &dwRead, NULL)
      && dwRead == sizeof(h)) {
    if (memcmp(h.signature, SIG, sizeof(SIG)) == 0) {
      pVolumeBitmap = (VOLUME_BITMAP_BUFFER *)
          LocalAlloc(LMEM_FIXED, h.uiBitmapSize);
      if (!ReadFile(hFile, pVolumeBitmap, h.uiBitmapSize,
          &dwRead, NULL) || dwRead != h.uiBitmapSize) {
        wprintf(L"Backup corrupted - could not read the "
                L"volume bitmap!\n");
        return ERROR_DATA_NOT_ACCEPTED;
      }
    } else
      SetFilePointer(hFile, 0, NULL, FILE_BEGIN);
  }

  DWORD dwErr = Transfer(hFile,
                         hDev,
                         hDev,
                         &h,
                         pVolumeBitmap,
                         buffer,
                         cbBufferSize,
                         ui64BytesWritten);

  DeviceIoControl(hDev,
                  IOCTL_DISK_UPDATE_PROPERTIES,
                  NULL,
                  0,
                  NULL,
                  0,
                  &dwRead,
                  NULL);

  return dwErr;
}

int _tmain(int argc, WCHAR **argv) {
  if (argc < 4) {
    wprintf(L"Usage: %s {disk|volume} cmd "
            L"args\n", argv[0]);
    wprintf(L"Where disk is physical number "
            L"of the drive, e. g. 0\n");
    wprintf(L"Volume is a drive letter or a "
            L"full path to \n");
    wprintf(L"the mount point, e. g. a: or "
            L"c:\\mount\\vol\n");
    wprintf(L"Commands:\n");
    wprintf(L"readto filename - read the "
            L"contents of the device into "
            L"a file\n");
    wprintf(L"writefrom filename - write the "
            L"contents of the file onto "
            L"the disk or volume\n");
    return 0;
  }

  UINT uiErrModeSav = SetErrorMode(
      SEM_FAILCRITICALERRORS);

  WCHAR szDevName[_MAX_PATH];
  if (!GetDeviceName(szDevName,
                     ARRLEN(szDevName),
                     argv[1])) {
    wprintf(L"Could not recognize %s\n", argv[1]);
    SetErrorMode(uiErrModeSav);
    return 1;
  }

  const int kBufferSize = 64 * 1024;
  void *buffer = VirtualAlloc(
      NULL,
      kBufferSize,
      MEM_COMMIT | MEM_RESERVE,
      PAGE_READWRITE);
  if (!buffer) {
    DWORD dwErr = GetLastError();
    wprintf (L"Could not reserve memory, "
             L"error %d (0x%08x)\n", dwErr, dwErr);
    SetErrorMode(uiErrModeSav);
    return 3;
  }

  DWORD dwTickStart = GetTickCount();
  unsigned __int64 ui64Bytes = 0;
  DWORD dwErr = ERROR_SUCCESS;
  if (_wcsicmp(argv[2], L"readto") == 0) {
    HANDLE hDev = OpenBlockDevice(szDevName, false);
    if (hDev != INVALID_HANDLE_VALUE) {
      HANDLE hFile = CreateFile(
          argv[3], GENERIC_WRITE, 0, NULL,
          CREATE_ALWAYS, 0, NULL);
      if (hFile != INVALID_HANDLE_VALUE) {
        wprintf(L"%s --> %s:\n", szDevName, argv[3]);
        dwErr = ReadToFile(hDev,
                           hFile,
                           buffer,
                           kBufferSize,
                           ui64Bytes);
        CloseHandle(hFile);
        if (dwErr != ERROR_SUCCESS)
          DeleteFile(argv[3]);
      } else {
        dwErr = GetLastError();
        wprintf (L"Could not open %s, error %d "
                 L"(0x%08x)\n", argv[3],
                 dwErr, dwErr);
      }
      CloseHandle(hDev);
    } else {
      dwErr = GetLastError();
      wprintf(L"Could not open %s, "
              L"error %d (0x%08x)\n",
              argv[1], dwErr, dwErr);
    }
  } else if (_wcsicmp(argv[2], L"writefrom") == 0) {
    HANDLE hDev = OpenBlockDevice(szDevName, true);
    if (hDev != INVALID_HANDLE_VALUE) {
      HANDLE hFile = CreateFile(
          argv[3], GENERIC_READ, 0, NULL,
          OPEN_EXISTING, 0, NULL);
      if (hFile != INVALID_HANDLE_VALUE) {
        wprintf(L"%s --> %s:\n", argv[3], szDevName);
        dwErr = WriteFromFile(hDev,
                              hFile,
                              buffer,
                              kBufferSize,
                              ui64Bytes);
        CloseHandle(hFile);
      } else {
        dwErr = GetLastError();
        wprintf (L"Could not open %s, error %d "
                 L"(0x%08x)\n", argv[3],
                 dwErr, dwErr);
      }
      CloseHandle(hDev);
    } else {
      dwErr = GetLastError();
      wprintf(L"Could not open %s, "
              L"error %d (0x%08x)\n",
              argv[1], dwErr, dwErr);
    }
  } else {
    wprintf (L"Command %s not recognized.\n",
             argv[2]);
    dwErr = ERROR_NOT_SUPPORTED;
  }

  if (dwErr == ERROR_SUCCESS &&
      ui64Bytes != 0) {
    wprintf(L"Transferred %I64u bytes in %d "
            L"seconds.\n", ui64Bytes,
            (GetTickCount() - dwTickStart) / 1000);
  }

  VirtualFree(buffer, 0, MEM_RELEASE);
  SetErrorMode(uiErrModeSav);

  return dwErr == ERROR_SUCCESS ? 0 : 4;
}