{*******************************************************}
{                                                       }
{       Tocsg.Disk                                      }
{                                                       }
{       Copyright (C) 2022 kkuzil                       }
{                                                       }
{*******************************************************}

unit Tocsg.Disk;

interface

uses
  System.SysUtils, Winapi.Windows, System.Generics.Collections;

const
  DISKPART_FMT_SELECT = 'select disk %d';
  DISKPART_SET_ATTRIBUTES_READONLY = 'attributes disk set readonly';
  DISKPART_DEL_ATTRIBUTES_READONLY = 'attributes disk clear readonly';
  DISKPART_OFFLINE_DISK = 'offline disk';
  DISKPART_ONLINE_DISK = 'online disk';

type
  PDriveExtent = ^TDriveExtent;
  TDriveExtent = record
    dwDiskNumber     : DWORD;
    liStartingOffset : LARGE_INTEGER;
    liExtentLength   : LARGE_INTEGER;
  end;

  PDiskExtents = ^TDiskExtents;
  TDiskExtents = record
    dwNumberOfDriveExtents: DWORD;
    Extents: array [0..0] of TDriveExtent;
  end;

  PDriveExtentEntry = ^TDriveExtentEntry;
  TDriveExtentEntry = record
    wDiskNum: WORD;
    sDrive: String;
    ullStartSector,
    ullTotalSectors: ULONGLONG;
  end;
  TDriveEntentList = TList<PDriveExtentEntry>;

  TTgDriveExtent = class(TObject)
  private
    lstEntry_: TDriveEntentList;
    procedure OnEntryNotify(Sender: TObject; const Item: PDriveExtentEntry;
                                Action: TCollectionNotification);
    function GetCount: Integer;
    function GetItem(nIndex: Integer): PDriveExtentEntry;
  public
    Constructor Create;
    Destructor Destroy; override;
    procedure RefreshDriveExtent;

    property Count: Integer read GetCount;
    property Items[nIndex: Integer]: PDriveExtentEntry read GetItem;
  end;

  PDevPathLetter = ^TDevPathLetter;
  TDevPathLetter = record
    cLetter: Char;
    sDevicePath: String;
  end;
  TDevPathLetterList = class(TList<PDevPathLetter>)
  protected
    procedure Notify(const Item: PDevPathLetter; Action: TCollectionNotification); override;
  end;

function GetDriveFromMask(dwMask: DWORD): String;
function GetDrivesFromMask(dwMask: DWORD; bLetterOnly: Boolean = false; bFixedOnly: Boolean = false): String;
function GetDrivesDevPathLetterList(aList: TDevPathLetterList): Integer;
function GetDrivesVolNumLetterList(aList: TDevPathLetterList): Integer;

function GetDriveExtent(const sDrive: String): TDriveExtent;

function GetDriveSize(const sDrive: String): ULONGLONG;
function GetDriveFree(const sDrive: String): ULONGLONG;

function GetVolumeSerial(const sDrive: String): DWORD;
function GetVolumeName(const sDrive: String): String;
function GetVolumeFilesystem(const sDrive: String): String;

function GetDriveTypeToStr(nType: Integer): String;
function IsReadOnlyByWriteProbe(const sDrive: String): Integer;
procedure SetDriveReadOnly(sDrive: String; nDiskNum: Integer; bVal: Boolean);

implementation

uses
  Tocsg.Strings, Tocsg.Exception, System.Classes, Tocsg.Path, Tocsg.Safe,
  Tocsg.Process, Tocsg.Files;

function GetDriveFromMask(dwMask: DWORD): String;
var
  wLetter: Word;
begin
  wLetter := Ord('A');
  while (dwMask and 1) = 0 do
  begin
    dwMask := dwMask shr 1;
    Inc(wLetter);
  end;
  Result := Format('%s:\', [Char(wLetter)]);
end;

function GetDrivesFromMask(dwMask: DWORD; bLetterOnly: Boolean = false; bFixedOnly: Boolean = false): String;
var
  ucDrive: Byte;
  sDrive: String;
begin
  Result := '';
  for ucDrive := 0 to 31 do
    if (dwMask and (1 shl ucDrive)) > 0 then
      if bLetterOnly then
      begin
        if bFixedOnly then
        begin
          sDrive := Format('%s:\', [Char(ucDrive + 65)]);
          if (GetDriveType(PChar(sDrive)) <> DRIVE_FIXED) then
            continue;
        end;

        SumString(Result, Char(ucDrive + 65), ',');
      end else begin
        sDrive := Format('%s:\', [Char(ucDrive + 65)]);
        if bFixedOnly and (GetDriveType(PChar(sDrive)) <> DRIVE_FIXED) then
          continue;

        SumString(Result, sDrive, ',');
      end;
end;

function GetDrivesDevPathLetterList(aList: TDevPathLetterList): Integer;
var
  dwMask: DWORD;
  ucDrive: Byte;
  cDrive: Char;
  pEnt: PDevPathLetter;
  arrPath: array [0..MAX_PATH] of Char;
begin
  try
    Result := 0;
    dwMask := GetLogicalDrives;
    for ucDrive := 0 to 31 do
      if (dwMask and (1 shl ucDrive)) > 0 then
      begin
        cDrive := Char(ucDrive + 65);
        if GetDriveExtent(cDrive).liExtentLength.QuadPart <> 0 then
        begin
          ZeroMemory(@arrPath, SizeOf(arrPath));
          QueryDosDevice(PChar(cDrive + ':'), @arrPath, MAX_PATH);
          if arrPath <> '' then
          begin
            New(pEnt);
            pEnt.cLetter := cDrive;
            pEnt.sDevicePath := arrPath;
            aList.Add(pEnt);
            Inc(Result);
          end;
        end;
      end;
  except
    on E: Exception do
      ETgException.TraceException(E, 'Fail .. GetDrivesDevPathLetterList()');
  end;
end;

function GetDrivesVolNumLetterList(aList: TDevPathLetterList): Integer;
var
  dwMask: DWORD;
  ucDrive: Byte;
  cDrive: Char;
  pEnt: PDevPathLetter;
  sVolNum: String;
begin
  Result := 0;
  dwMask := GetLogicalDrives;
  for ucDrive := 0 to 31 do
    if (dwMask and (1 shl ucDrive)) > 0 then
    begin
      cDrive := Char(ucDrive + 65);
      if GetDriveExtent(cDrive).liExtentLength.QuadPart <> 0 then
      begin
        sVolNum := IntToHex(GetVolumeSerial(cDrive + ':\'), 3);
        if sVolNum <> '' then
        begin
          New(pEnt);
          pEnt.cLetter := cDrive;
          pEnt.sDevicePath := sVolNum;
          aList.Add(pEnt);
          Inc(Result);
        end;
      end;
    end;
end;

function GetDriveExtent(const sDrive: String): TDriveExtent;
var
  hVolume: THandle;
  DiskExtents: PDiskExtents;
  dwOutBytes: DWORD;
  sDrivePath: String;
begin
  ZeroMemory(@Result, SizeOf(Result));
  Result.dwDiskNumber := 999;
  sDrivePath := Format('\\.\%s:', [sDrive[1]]);

  hVolume := CreateFile(PChar(sDrivePath),
                        0,//GENERIC_READ or GENERIC_WRITE, // GENERIC_READ 이걸 넣어주면 딜레이가 많이 생긴다.. 15_0202 16:26:50 sunk
                        FILE_SHARE_READ,// or FILE_SHARE_WRITE,
                        nil,
                        OPEN_EXISTING,
                        0,
                        0);

//  hVolume := CreateFile(PChar(sDrivePath),
//                        GENERIC_READ, // or GENERIC_WRITE,
//                        FILE_SHARE_READ,// or FILE_SHARE_WRITE,
//                        nil,
//                        OPEN_EXISTING,
//                        FILE_FLAG_NO_BUFFERING or FILE_FLAG_SEQUENTIAL_SCAN,
//                        0);

  if hVolume < 1 then
    exit;

  DiskExtents := AllocMem(MAX_PATH);
  try
    if DeviceIoControl(hVolume,
                       IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS,
                       nil,
                       0,
                       DiskExtents,
                       MAX_PATH,
                       dwOutBytes,
                       nil)  then
    begin
      if DiskExtents.dwNumberOfDriveExtents > 0 then
        Result := DiskExtents.Extents[0];
    end;
//    else // imdisk로 마운트한 드라이브 정보를 가져오기 위해서 아래처럼 시도 하지만.. 정보를 제대로 가져오지 못하는 상황이다 15_0202 17:45:02 sunk
//    if DeviceIoControl(hVolume,
//                       IOCTL_DISK_GET_DRIVE_GEOMETRY,
//                       nil,
//                       0,
//                       @DriveExtent,
//                       SizeOf(DriveExtent),
//                       dwOutBytes,
//                       nil)  then
//    begin
//      Result := DriveExtent;
//    end;

  // todo : IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS 이걸로 안되는거 IOCTL_DISK_GET_DRIVE_GEOMETRY_EX 이걸로 한번더 확인해 보기

  finally
    Finalize(DiskExtents);
    FreeMem(DiskExtents, Max_Path);
    CloseHandle(hVolume);
  end;
end;

function GetDriveSize(const sDrive: String): ULONGLONG;
var
  bResult: Boolean;
  llAvail,
  llTotal,
  llFree: TLargeInteger;
begin
  Result := 0;

// 윈도우 xp에서 빈디스크 에러 메시지 걸러내기 2012-01-25 sunk
// 보안, 가상 드라이브의 경우 IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS가 항상 실패 25_0423 08:45:20 kku
//  if GetDriveExtent(sDrive[1]).liExtentLength.QuadPart = 0 then
//    exit;

  bResult := GetDiskFreeSpaceEx(PChar(sDrive), llAvail, llTotal, @llFree);
  if bResult = true then
    Result := llTotal;
end;

function GetDriveFree(const sDrive: String): ULONGLONG;
var
  FSectorsPerCluster,
  FBytesPerSector,
  FNumberOfFreeClusters,

  FTotalNumberOfClusters: DWORD;
begin
  Result := 0;
  if sDrive = '' then
    exit;

  if GetDriveExtent(sDrive[1]).liExtentLength.QuadPart = 0 then
    exit;

  FSectorsPerCluster := 0;
  FBytesPerSector := 0;
  FNumberOfFreeClusters := 0;
  FTotalNumberOfClusters := 0;

  // GetDiskFreeSpace() API 호출
  GetDiskFreeSpace(
                    PChar(sDrive),
                    FSectorsPerCluster,
                    FBytesPerSector,
                    FNumberOfFreeClusters,
                    FTotalNumberOfClusters
                   );
  Result := ULONGLONG(FBytesPerSector) * FSectorsPerCluster * FNumberOfFreeClusters;
end;

function GetVolumeSerial(const sDrive: String): DWORD;
var
  sVolumeName: array [0..255] of Char;
  dwMaxComponentLen: DWORD;
  sFileSystemName: array[0..255] of Char;
  dwFileSystemFlag: DWORD;
  bResult: Boolean;
//  ErrCode: DWORD;
begin
  Result := 0;
  bResult := GetVolumeInformation(PWideChar(sDrive), sVolumeName, 256,
                       @Result, dwMaxComponentLen, dwFileSystemFlag,
                       sFileSystemName, 256);

  if not bResult then
  begin
//    ErrCode := GetLastError;
//    OutputDebugString(PChar('GetVolumeInformation Fail!! Error Code : ' + IntToStr(ErrCode)));
  end;
end;

function GetVolumeName(const sDrive: String): String;
var
  sVolumeName: array [0..255] of Char;
  dwVolumeSerial: DWORD;
  dwMaxComponentLen: DWORD;
  sFileSystemName: array[0..255] of Char;
  dwFileSystemFlag: DWORD;
begin
  Result := '';
  if GetVolumeInformation(PChar(sDrive), sVolumeName, 256,
                       @dwVolumeSerial, dwMaxComponentLen, dwFileSystemFlag,
                       sFileSystemName, 256) then Result := sVolumeName;
end;

function GetVolumeFilesystem(const sDrive: String): String;
var
  sVolumeName: array [0..255] of Char;
  dwVolumeSerial: DWORD;
  dwMaxComponentLen: DWORD;
  sFileSystemName: array[0..255] of Char;
  dwFileSystemFlag: DWORD;
begin
  Result := '';
  if GetVolumeInformation(PChar(sDrive), sVolumeName, 256,
                       @dwVolumeSerial, dwMaxComponentLen, dwFileSystemFlag,
                       sFileSystemName, 256) then Result := sFileSystemName;
end;

function GetDriveTypeToStr(nType: Integer): String;
begin
  case nType of
    DRIVE_FIXED     : Result := 'FIXED';
    DRIVE_CDROM     : Result := 'CDROM';
    DRIVE_REMOVABLE : Result := 'REMOVABLE';
    else Result := 'Unknown';
  end;
end;

// 0 : 쓰기 가능, 1 : 읽기 전용, 2 : 권한 없음, 3 : 공간 부족
function IsReadOnlyByWriteProbe(const sDrive: String): Integer;
var
  sTmpPath: string;
  h: THandle;
begin
  Result := 0;
  // 루트에 직접 파일 쓰기 시도 (권한 문제를 피하려면 하위 폴더를 탐색해서 시도해도 됨)
  sTmpPath := IncludeTrailingPathDelimiter(sDrive) + '_wr_probe_.tmp';

  if FileExists(sTmpPath) then
    DeleteFile(PChar(sTmpPath));

  h := CreateFile(PChar(sTmpPath),
                  GENERIC_WRITE,
                  FILE_SHARE_READ or FILE_SHARE_WRITE,
                  nil,
                  CREATE_NEW,
                  FILE_ATTRIBUTE_TEMPORARY or FILE_FLAG_DELETE_ON_CLOSE,
                  0);
  if h = INVALID_HANDLE_VALUE then
  begin
    case GetLastError of
      ERROR_WRITE_PROTECT : Result := 1;
      ERROR_ACCESS_DENIED : Result := 2;
      ERROR_DISK_FULL     : Result := 3;
    end
  end else
    CloseHandle(h);
end;

// ThdUsbMon.pas 에서 SetReadOnly() 을 가져옴 25_1015 16:21:30 kku
procedure SetDriveReadOnly(sDrive: String; nDiskNum: Integer; bVal: Boolean);
var
  StrList: TStringList;
  sTemp,
  sDrvPath,
  sScptPath: String;
  nTO,
  nRst: Integer;
  fs: TFileStream;
begin
  if sDrive = '' then
    exit;

  if nDiskNum <= 0 then
    exit;

  sScptPath := GetRunExePathDir + Format('$d-scrpt=%d.txt', [GetTickCount]);
  Guard(StrList, TStringList.Create);
  StrList.Add(Format(DISKPART_FMT_SELECT, [nDiskNum]));
  if bVal then
    StrList.Add(DISKPART_SET_ATTRIBUTES_READONLY)
  else
    StrList.Add(DISKPART_DEL_ATTRIBUTES_READONLY);
  StrList.SaveToFile(sScptPath, TEncoding.ANSI);
  ExecuteAppWaitUntilTerminate('diskpart.exe', Format('/s "%s"', [sScptPath]), SW_HIDE, 10000);
  DeleteFile_wait(PChar(sScptPath), 2);
end;

{ TTgDriveExtent }

Constructor TTgDriveExtent.Create;
begin
  Inherited Create;
  lstEntry_ := TDriveEntentList.Create;
  lstEntry_.OnNotify := OnEntryNotify;
  RefreshDriveExtent;
end;

Destructor TTgDriveExtent.Destroy;
begin
  FreeAndNil(lstEntry_);
  Inherited;
end;

procedure TTgDriveExtent.OnEntryNotify(Sender: TObject; const Item: PDriveExtentEntry;
                               Action: TCollectionNotification);
begin
  case Action of
    cnAdded: ;
    cnRemoved: Dispose(Item);
    cnExtracted: ;
  end;
end;

function TTgDriveExtent.GetCount: Integer;
begin
  Result := lstEntry_.Count;
end;

function TTgDriveExtent.GetItem(nIndex: Integer): PDriveExtentEntry;
begin
  Result := nil;
  if (nIndex > -1) and (nIndex < lstEntry_.Count) then
    Result := lstEntry_[nIndex];
end;

procedure TTgDriveExtent.RefreshDriveExtent;
var
  dwDriveMask: DWORD;
  ucDrive: Byte;
  sDrive: String;
  Extent: TDriveExtent;
  pEntry: PDriveExtentEntry;
begin
  lstEntry_.Clear;

  dwDriveMask := GetLogicalDrives;
  for ucDrive := 0 to 31 do
    if (dwDriveMask and (1 shl ucDrive)) > 0 then
    begin
      sDrive := Format('%s:\', [Char(ucDrive + 65)]);
      Extent := GetDriveExtent(sDrive);
      if Extent.liExtentLength.QuadPart > 0 then
      begin
        case GetDriveType(PChar(sDrive)) of
          DRIVE_UNKNOWN,
          DRIVE_REMOTE,
          DRIVE_CDROM: continue;
        end;

        New(pEntry);
        ZeroMemory(pEntry, SizeOf(TDriveExtentEntry));
        pEntry.sDrive := sDrive;
        with Extent do
        begin
          pEntry.wDiskNum := dwDiskNumber;
          pEntry.ullStartSector := liStartingOffset.QuadPart div 512;
          pEntry.ullTotalSectors := liExtentLength.QuadPart div 512;
        end;
        lstEntry_.Add(pEntry);
      end;
    end;
end;

{ TDevPathLetterList }

procedure TDevPathLetterList.Notify(const Item: PDevPathLetter; Action: TCollectionNotification);
begin
  case Action of
    cnAdded: ;
    cnRemoved: Dispose(Item);
    cnExtracted: ;
  end;
end;

end.