[DISKPART] Implement the clean command

This commit is contained in:
Eric Kohl 2022-06-17 01:37:26 +02:00
parent 74870c1ae1
commit 2be7af18fe
14 changed files with 390 additions and 10 deletions

View file

@ -8,7 +8,241 @@
#include "diskpart.h"
BOOL clean_main(INT argc, LPWSTR *argv)
#define NDEBUG
#include <debug.h>
BOOL
clean_main(
_In_ INT argc,
_In_ PWSTR *argv)
{
PLIST_ENTRY Entry;
PPARTENTRY PartEntry;
PVOLENTRY VolumeEntry;
BOOL bAll = FALSE;
PUCHAR SectorsBuffer = NULL;
ULONG LayoutBufferSize, Size;
INT i;
WCHAR Buffer[MAX_PATH];
UNICODE_STRING Name;
OBJECT_ATTRIBUTES ObjectAttributes;
IO_STATUS_BLOCK IoStatusBlock;
HANDLE FileHandle = NULL;
LARGE_INTEGER Offset, Count, MaxCount;
NTSTATUS Status;
DPRINT("Clean()\n");
if (CurrentDisk == NULL)
{
ConResPuts(StdOut, IDS_SELECT_NO_DISK);
return TRUE;
}
/* Do not allow to clean the boot disk */
if ((CurrentDisk->BiosFound == TRUE) &&
(CurrentDisk->BiosDiskNumber == 0))
{
ConResPuts(StdOut, IDS_CLEAN_SYSTEM);
return TRUE;
}
for (i = 1; i < argc; i++)
{
if (_wcsicmp(argv[1], L"all") == 0)
{
bAll = TRUE;
}
}
/* Dismount and remove all logical partitions */
while (!IsListEmpty(&CurrentDisk->LogicalPartListHead))
{
Entry = RemoveHeadList(&CurrentDisk->LogicalPartListHead);
PartEntry = CONTAINING_RECORD(Entry, PARTENTRY, ListEntry);
/* Dismount the logical partition */
if (PartEntry->PartitionType != 0)
{
DismountVolume(PartEntry);
VolumeEntry = GetVolumeFromPartition(PartEntry);
if (VolumeEntry)
RemoveVolume(VolumeEntry);
}
/* Delete it */
RtlFreeHeap(RtlGetProcessHeap(), 0, PartEntry);
}
/* Dismount and remove all primary partitions */
while (!IsListEmpty(&CurrentDisk->PrimaryPartListHead))
{
Entry = RemoveHeadList(&CurrentDisk->PrimaryPartListHead);
PartEntry = CONTAINING_RECORD(Entry, PARTENTRY, ListEntry);
/* Dismount the primary partition */
if ((PartEntry->PartitionType != 0) &&
(IsContainerPartition(PartEntry->PartitionType) == FALSE))
{
DismountVolume(PartEntry);
VolumeEntry = GetVolumeFromPartition(PartEntry);
if (VolumeEntry)
RemoveVolume(VolumeEntry);
}
/* Delete it */
RtlFreeHeap(RtlGetProcessHeap(), 0, PartEntry);
}
/* Initialize the disk entry */
CurrentDisk->ExtendedPartition = NULL;
CurrentDisk->Dirty = FALSE;
CurrentDisk->NewDisk = TRUE;
CurrentDisk->NoMbr = TRUE;
/* Wipe the layout buffer */
RtlFreeHeap(RtlGetProcessHeap(), 0, CurrentDisk->LayoutBuffer);
LayoutBufferSize = sizeof(DRIVE_LAYOUT_INFORMATION) +
((4 - ANYSIZE_ARRAY) * sizeof(PARTITION_INFORMATION));
CurrentDisk->LayoutBuffer = RtlAllocateHeap(RtlGetProcessHeap(),
HEAP_ZERO_MEMORY,
LayoutBufferSize);
if (CurrentDisk->LayoutBuffer == NULL)
{
DPRINT1("Failed to allocate the disk layout buffer!\n");
return TRUE;
}
/* Allocate a 1MB sectors buffer */
SectorsBuffer = RtlAllocateHeap(RtlGetProcessHeap(),
HEAP_ZERO_MEMORY,
1024 * 1024);
if (SectorsBuffer == NULL)
{
DPRINT1("Failed to allocate the sectors buffer!\n");
goto done;
}
/* Open the disk for writing */
StringCchPrintfW(Buffer, ARRAYSIZE(Buffer),
L"\\Device\\Harddisk%d\\Partition0",
CurrentDisk->DiskNumber);
RtlInitUnicodeString(&Name, Buffer);
InitializeObjectAttributes(&ObjectAttributes,
&Name,
OBJ_CASE_INSENSITIVE,
NULL,
NULL);
Status = NtOpenFile(&FileHandle,
GENERIC_READ | GENERIC_WRITE | SYNCHRONIZE,
&ObjectAttributes,
&IoStatusBlock,
0,
FILE_SYNCHRONOUS_IO_NONALERT);
if (!NT_SUCCESS(Status))
{
DPRINT1("Failed to open the disk! (Status 0x%08lx)\n", Status);
ConResPuts(StdOut, IDS_CLEAN_FAIL);
goto done;
}
/* Clean sectors */
if (bAll)
{
MaxCount.QuadPart = (CurrentDisk->SectorCount.QuadPart * CurrentDisk->BytesPerSector) / (1024 * 1024);
for (Count.QuadPart = 0; Count.QuadPart < MaxCount.QuadPart; Count.QuadPart++)
{
Offset.QuadPart = Count.QuadPart * (1024 * 1024);
Status = NtWriteFile(FileHandle,
NULL,
NULL,
NULL,
&IoStatusBlock,
SectorsBuffer,
1024 * 1024,
&Offset,
NULL);
if (!NT_SUCCESS(Status))
{
DPRINT1("Failed to write MB! (Status 0x%08lx)\n", Status);
ConResPuts(StdOut, IDS_CLEAN_FAIL);
goto done;
}
}
Size = (ULONG)(CurrentDisk->SectorCount.QuadPart * CurrentDisk->BytesPerSector) % (1024 * 1024);
if (Size != 0)
{
Offset.QuadPart += (1024 * 1024);
Status = NtWriteFile(FileHandle,
NULL,
NULL,
NULL,
&IoStatusBlock,
SectorsBuffer,
Size,
&Offset,
NULL);
if (!NT_SUCCESS(Status))
{
DPRINT1("Failed to write the last part! (Status 0x%08lx)\n", Status);
ConResPuts(StdOut, IDS_CLEAN_FAIL);
goto done;
}
}
}
else
{
/* Clean the first MB */
Offset.QuadPart = 0;
Status = NtWriteFile(FileHandle,
NULL,
NULL,
NULL,
&IoStatusBlock,
SectorsBuffer,
1024 * 1024,
&Offset,
NULL);
if (!NT_SUCCESS(Status))
{
DPRINT1("Failed to write the first MB! (Status 0x%08lx)\n", Status);
ConResPuts(StdOut, IDS_CLEAN_FAIL);
goto done;
}
/* Clean the last MB */
Offset.QuadPart = (CurrentDisk->SectorCount.QuadPart * CurrentDisk->BytesPerSector) - (1024 * 1024);
Status = NtWriteFile(FileHandle,
NULL,
NULL,
NULL,
&IoStatusBlock,
SectorsBuffer,
1024 * 1024,
&Offset,
NULL);
if (!NT_SUCCESS(Status))
{
DPRINT1("Failed to write the last MB! (Status 0x%08lx)\n", Status);
ConResPuts(StdOut, IDS_CLEAN_FAIL);
goto done;
}
}
ConResPuts(StdOut, IDS_CLEAN_SUCCESS);
done:
if (FileHandle != NULL)
NtClose(FileHandle);
if (SectorsBuffer != NULL)
RtlFreeHeap(RtlGetProcessHeap(), 0, SectorsBuffer);
return TRUE;
}

View file

@ -22,6 +22,7 @@
#include <winioctl.h>
#include <errno.h>
#include <strsafe.h>
#include <conutils.h>
@ -439,7 +440,16 @@ GetPrimaryPartitionCount(
NTSTATUS
DismountVolume(
IN PPARTENTRY PartEntry);
_In_ PPARTENTRY PartEntry);
PVOLENTRY
GetVolumeFromPartition(
_In_ PPARTENTRY PartEntry);
VOID
RemoveVolume(
_In_ PVOLENTRY VolumeEntry);
/* recover.c */
BOOL recover_main(INT argc, LPWSTR *argv);

View file

@ -22,6 +22,13 @@ BEGIN
IDS_ACTIVE_ALREADY "\nDie aktuelle Partition wurde bereits als aktiv markiert.\n"
END
STRINGTABLE
BEGIN
IDS_CLEAN_FAIL "\nDer Datenträger konnte nicht bereinigt werden.\nMöglicherweise können die Daten auf diesem Datenträger nicht wiederhergestellt werden.\n"
IDS_CLEAN_SUCCESS "\nDer Datenträger wurde bereinigt.\n"
IDS_CLEAN_SYSTEM "\nDer gewählte Datenträger ist zum Ausführen des Computers erforderlich und kann nicht bereinigt werden.\n"
END
STRINGTABLE
BEGIN
IDS_CREATE_PARTITION_FAIL "\nDie angegebene Partition konnte nicht erstellt werden.\n"

View file

@ -22,6 +22,13 @@ BEGIN
IDS_ACTIVE_ALREADY "\nThe current partition is already marked as active.\n"
END
STRINGTABLE
BEGIN
IDS_CLEAN_FAIL "\nDiskPart was unable to clean the disk.\nThe data on this disk may be unrecoverable.\n"
IDS_CLEAN_SUCCESS "\nDiskPart succeeded in cleaning the disk.\n"
IDS_CLEAN_SYSTEM "\nThe selected disk is neccessary to the operation of your computer, and may not be cleaned.\n"
END
STRINGTABLE
BEGIN
IDS_CREATE_PARTITION_FAIL "\nDiskPart was unable to create the specified partition.\n"

View file

@ -22,6 +22,13 @@ BEGIN
IDS_ACTIVE_ALREADY "\nThe current partition is already marked as active.\n"
END
STRINGTABLE
BEGIN
IDS_CLEAN_FAIL "\nDiskPart was unable to clean the disk.\nThe data on this disk may be unrecoverable.\n"
IDS_CLEAN_SUCCESS "\nDiskPart succeeded in cleaning the disk.\n"
IDS_CLEAN_SYSTEM "\nThe selected disk is neccessary to the operation of your computer, and may not be cleaned.\n"
END
STRINGTABLE
BEGIN
IDS_CREATE_PARTITION_FAIL "\nDiskPart was unable to create the specified partition.\n"

View file

@ -24,6 +24,13 @@ BEGIN
IDS_ACTIVE_ALREADY "\nThe current partition is already marked as active.\n"
END
STRINGTABLE
BEGIN
IDS_CLEAN_FAIL "\nDiskPart was unable to clean the disk.\nThe data on this disk may be unrecoverable.\n"
IDS_CLEAN_SUCCESS "\nDiskPart succeeded in cleaning the disk.\n"
IDS_CLEAN_SYSTEM "\nThe selected disk is neccessary to the operation of your computer, and may not be cleaned.\n"
END
STRINGTABLE
BEGIN
IDS_CREATE_PARTITION_FAIL "\nDiskPart was unable to create the specified partition.\n"

View file

@ -24,6 +24,13 @@ BEGIN
IDS_ACTIVE_ALREADY "\nThe current partition is already marked as active.\n"
END
STRINGTABLE
BEGIN
IDS_CLEAN_FAIL "\nDiskPart was unable to clean the disk.\nThe data on this disk may be unrecoverable.\n"
IDS_CLEAN_SUCCESS "\nDiskPart succeeded in cleaning the disk.\n"
IDS_CLEAN_SYSTEM "\nThe selected disk is neccessary to the operation of your computer, and may not be cleaned.\n"
END
STRINGTABLE
BEGIN
IDS_CREATE_PARTITION_FAIL "\nDiskPart was unable to create the specified partition.\n"

View file

@ -24,6 +24,13 @@ BEGIN
IDS_ACTIVE_ALREADY "\nThe current partition is already marked as active.\n"
END
STRINGTABLE
BEGIN
IDS_CLEAN_FAIL "\nDiskPart was unable to clean the disk.\nThe data on this disk may be unrecoverable.\n"
IDS_CLEAN_SUCCESS "\nDiskPart succeeded in cleaning the disk.\n"
IDS_CLEAN_SYSTEM "\nThe selected disk is neccessary to the operation of your computer, and may not be cleaned.\n"
END
STRINGTABLE
BEGIN
IDS_CREATE_PARTITION_FAIL "\nDiskPart was unable to create the specified partition.\n"

View file

@ -26,6 +26,13 @@ BEGIN
IDS_ACTIVE_ALREADY "\nThe current partition is already marked as active.\n"
END
STRINGTABLE
BEGIN
IDS_CLEAN_FAIL "\nDiskPart was unable to clean the disk.\nThe data on this disk may be unrecoverable.\n"
IDS_CLEAN_SUCCESS "\nDiskPart succeeded in cleaning the disk.\n"
IDS_CLEAN_SYSTEM "\nThe selected disk is neccessary to the operation of your computer, and may not be cleaned.\n"
END
STRINGTABLE
BEGIN
IDS_CREATE_PARTITION_FAIL "\nDiskPart was unable to create the specified partition.\n"

View file

@ -24,6 +24,13 @@ BEGIN
IDS_ACTIVE_ALREADY "\nThe current partition is already marked as active.\n"
END
STRINGTABLE
BEGIN
IDS_CLEAN_FAIL "\nDiskPart was unable to clean the disk.\nThe data on this disk may be unrecoverable.\n"
IDS_CLEAN_SUCCESS "\nDiskPart succeeded in cleaning the disk.\n"
IDS_CLEAN_SYSTEM "\nThe selected disk is neccessary to the operation of your computer, and may not be cleaned.\n"
END
STRINGTABLE
BEGIN
IDS_CREATE_PARTITION_FAIL "\nDiskPart was unable to create the specified partition.\n"

View file

@ -31,6 +31,13 @@ BEGIN
IDS_ACTIVE_ALREADY "\nThe current partition is already marked as active.\n"
END
STRINGTABLE
BEGIN
IDS_CLEAN_FAIL "\nDiskPart was unable to clean the disk.\nThe data on this disk may be unrecoverable.\n"
IDS_CLEAN_SUCCESS "\nDiskPart succeeded in cleaning the disk.\n"
IDS_CLEAN_SYSTEM "\nThe selected disk is neccessary to the operation of your computer, and may not be cleaned.\n"
END
STRINGTABLE
BEGIN
IDS_CREATE_PARTITION_FAIL "\nDiskPart was unable to create the specified partition.\n"

View file

@ -25,6 +25,13 @@ BEGIN
IDS_ACTIVE_ALREADY "\nThe current partition is already marked as active.\n"
END
STRINGTABLE
BEGIN
IDS_CLEAN_FAIL "\nDiskPart was unable to clean the disk.\nThe data on this disk may be unrecoverable.\n"
IDS_CLEAN_SUCCESS "\nDiskPart succeeded in cleaning the disk.\n"
IDS_CLEAN_SYSTEM "\nThe selected disk is neccessary to the operation of your computer, and may not be cleaned.\n"
END
STRINGTABLE
BEGIN
IDS_CREATE_PARTITION_FAIL "\nDiskPart was unable to create the specified partition.\n"

View file

@ -10,7 +10,6 @@
#include "diskpart.h"
#include <ntddscsi.h>
#include <strsafe.h>
#define NDEBUG
#include <debug.h>
@ -1635,7 +1634,7 @@ IsSamePrimaryLayoutEntry(
ULONG
GetPrimaryPartitionCount(
IN PDISKENTRY DiskEntry)
_In_ PDISKENTRY DiskEntry)
{
PLIST_ENTRY Entry;
PPARTENTRY PartEntry;
@ -1657,7 +1656,7 @@ GetPrimaryPartitionCount(
static
ULONG
GetLogicalPartitionCount(
IN PDISKENTRY DiskEntry)
_In_ PDISKENTRY DiskEntry)
{
PLIST_ENTRY ListEntry;
PPARTENTRY PartEntry;
@ -1679,7 +1678,7 @@ GetLogicalPartitionCount(
static
BOOLEAN
ReAllocateLayoutBuffer(
IN PDISKENTRY DiskEntry)
_In_ PDISKENTRY DiskEntry)
{
PDRIVE_LAYOUT_INFORMATION NewLayoutBuffer;
ULONG NewPartitionCount;
@ -1731,7 +1730,7 @@ ReAllocateLayoutBuffer(
VOID
UpdateDiskLayout(
IN PDISKENTRY DiskEntry)
_In_ PDISKENTRY DiskEntry)
{
PPARTITION_INFORMATION PartitionInfo;
PPARTITION_INFORMATION LinkInfo = NULL;
@ -1901,7 +1900,7 @@ UpdateDiskLayout(
PPARTENTRY
GetPrevUnpartitionedEntry(
IN PPARTENTRY PartEntry)
_In_ PPARTENTRY PartEntry)
{
PDISKENTRY DiskEntry = PartEntry->DiskEntry;
PPARTENTRY PrevPartEntry;
@ -1930,7 +1929,7 @@ GetPrevUnpartitionedEntry(
PPARTENTRY
GetNextUnpartitionedEntry(
IN PPARTENTRY PartEntry)
_In_ PPARTENTRY PartEntry)
{
PDISKENTRY DiskEntry = PartEntry->DiskEntry;
PPARTENTRY NextPartEntry;
@ -1958,7 +1957,7 @@ GetNextUnpartitionedEntry(
NTSTATUS
DismountVolume(
IN PPARTENTRY PartEntry)
_In_ PPARTENTRY PartEntry)
{
NTSTATUS Status;
NTSTATUS LockStatus;
@ -2063,4 +2062,67 @@ DismountVolume(
return Status;
}
PVOLENTRY
GetVolumeFromPartition(
_In_ PPARTENTRY PartEntry)
{
PLIST_ENTRY Entry;
PVOLENTRY VolumeEntry;
ULONG i;
if ((PartEntry == NULL) ||
(PartEntry->DiskEntry == NULL))
return NULL;
Entry = VolumeListHead.Flink;
while (Entry != &VolumeListHead)
{
VolumeEntry = CONTAINING_RECORD(Entry, VOLENTRY, ListEntry);
if (VolumeEntry->pExtents == NULL)
return NULL;
for (i = 0; i < VolumeEntry->pExtents->NumberOfDiskExtents; i++)
{
if (VolumeEntry->pExtents->Extents[i].DiskNumber == PartEntry->DiskEntry->DiskNumber)
{
if ((VolumeEntry->pExtents->Extents[i].StartingOffset.QuadPart == PartEntry->StartSector.QuadPart * PartEntry->DiskEntry->BytesPerSector) &&
(VolumeEntry->pExtents->Extents[i].ExtentLength.QuadPart == PartEntry->SectorCount.QuadPart * PartEntry->DiskEntry->BytesPerSector))
return VolumeEntry;
}
}
Entry = Entry->Flink;
}
return NULL;
}
VOID
RemoveVolume(
_In_ PVOLENTRY VolumeEntry)
{
if (VolumeEntry == NULL)
return;
if (VolumeEntry == CurrentVolume)
CurrentVolume = NULL;
RemoveEntryList(&VolumeEntry->ListEntry);
if (VolumeEntry->pszLabel)
RtlFreeHeap(RtlGetProcessHeap(), 0, VolumeEntry->pszLabel);
if (VolumeEntry->pszFilesystem)
RtlFreeHeap(RtlGetProcessHeap(), 0, VolumeEntry->pszFilesystem);
if (VolumeEntry->pExtents)
RtlFreeHeap(RtlGetProcessHeap(), 0, VolumeEntry->pExtents);
/* Release disk entry */
RtlFreeHeap(RtlGetProcessHeap(), 0, VolumeEntry);
}
/* EOF */

View file

@ -22,6 +22,10 @@
#define IDS_ACTIVE_SUCCESS 1001
#define IDS_ACTIVE_ALREADY 1002
#define IDS_CLEAN_FAIL 1020
#define IDS_CLEAN_SUCCESS 1021
#define IDS_CLEAN_SYSTEM 1022
#define IDS_CREATE_PARTITION_FAIL 1050
#define IDS_CREATE_PARTITION_SUCCESS 1051