{*******************************************************}
{                                                       }
{       ThdUsbMon                                       }
{                                                       }
{       Copyright (C) 2022 kku                          }
{                                                       }
{*******************************************************}

unit ThdUsbMon;

interface

uses
  Tocsg.Thread, System.SysUtils, System.Classes, Winapi.Windows,
  Winapi.Messages, System.Generics.Collections, Tocsg.Driver,
  Tocsg.USB;

const
  WM_USBCONTROL_NOTIFY = WM_USER + 8654;

  ACTION_USBCONTROL_ARRIVAL = 1;
  ACTION_USBCONTROL_REMOVE = 2;
  ACTION_USBCONTROL_READONLY = 3;
  ACTION_USBCONTROL_READONLY_RETRAY = 4;
  ACTION_USBCONTROL_BLOCK = 5;

type
  PUsbEnt = ^TUsbEnt;
  TUsbEnt = record
    DriveInfo: TDriveInfo;
    bReadOnly,
    bRoPopup: Boolean;
    dwRoTick: DWORD;
  end;

  TUsbProcKind = (upkMonitor, upkReadOnly, upkBlock);
  TThdUsbMon = class(TTgThread)
  protected
    hRcvWnd_: HWND;
    DriveList_: TList<PUsbEnt>;
    EmDriveList_: TStringList;
    qArrivalUsb_: TQueue<String>;
    ActionKind_: TUsbProcKind;
    UsbNoti_: TTgUSBEventNotify;
    sIgrUsbSerials_: String;
    IgrUsbSerialList_: TStringList;
    bRecoverWB_: Boolean;
    procedure OnUsbNotify(Sender: TObject; const Item: PUsbEnt; Action: TCollectionNotification);
    procedure OnUSBArrival(Sender: TObject; pInfo: PDevBroadcastVolume);
    procedure OnUSBQueryRemove(Sender: TObject; sDrive: String; var bAccept: Boolean);
    procedure OnUSBRemove(Sender: TObject; pInfo: PDevBroadcastVolume);
    function GetDriveEntByDrive(sDrive: String): PUsbEnt;
    procedure ProcessUsbControl(pEnt: PUsbEnt);
    function GetRecoverWB: Boolean;
    procedure UpdateDriveInfo;
    procedure Execute; override;
  public
    Constructor Create(hRcvWnd: HWND; aKind: TUsbProcKind);
    Destructor Destroy; override;

    procedure AddDrive(sDrive: String);
    procedure AddEmptyDrive(sDrive: String);
    procedure DelDrive(sDrive: String);
    procedure SetRecoverWB(bVal: Boolean);
  end;

procedure SetReadOnly(sDrive: String; nDiskNum: Integer; bVal: Boolean);

implementation

uses
{$IFDEF _HE_}
  ManagerService, GlobalDefine, Tocsg.Convert, ManagerModel, Condition,
{$ENDIF}
  Tocsg.Safe, Tocsg.Path, Tocsg.Shell, Tocsg.Process,
  Tocsg.Files, Tocsg.Strings, Tocsg.Exception, Tocsg.Disk;

procedure SetReadOnly(sDrive: String; nDiskNum: Integer; bVal: Boolean);
var
  StrList: TStringList;
  sTemp,
  sDrvPath,
  sScptPath: String;
  nTO,
  nRst: Integer;
  fs: TFileStream;
  Disk: TTgDriver;
begin
  if sDrive = '' then
    exit;

  if nDiskNum <= 0 then
    exit;

{$IFDEF _HE_}
  try
    if gMgSvc.FileService <> nil then
      gMgSvc.FileService.DelDriveWatch(sDrive);
  except
  // ..
  end;
{$ENDIF}

  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);

{$IFDEF _HE_}
  if not bVal then
  begin
    try
      if gMgSvc.FileService <> nil then
        gMgSvc.FileService.AddDriveWatch(sDrive, true);
    except
    // ..
    end;
  end;
{$ENDIF}
end;

{ TThdUsbMon }

Constructor TThdUsbMon.Create(hRcvWnd: HWND; aKind: TUsbProcKind);
begin
  Inherited Create;
  hRcvWnd_ := hRcvWnd;
  ActionKind_ := aKind;
  DriveList_ := TList<PUsbEnt>.Create;
  DriveList_.OnNotify := OnUsbNotify;
  EmDriveList_ := TStringList.Create;
  EmDriveList_.CaseSensitive := false;
  qArrivalUsb_ := TQueue<String>.Create;
  sIgrUsbSerials_ := '';
  IgrUsbSerialList_ := TStringList.Create;
  bRecoverWB_ := true;
{$IFDEF _HE_}
  UsbNoti_ := nil;
{$ELSE}
  UsbNoti_ := TTgUSBEventNotify.Create;
  UsbNoti_.OnUSBArrival := OnUSBArrival;
  UsbNoti_.OnUSBQueryRemove := OnUSBQueryRemove;
  UsbNoti_.OnUSBRemove := OnUSBRemove;
{$ENDIF}
end;

Destructor TThdUsbMon.Destroy;
var
  i: Integer;
begin
  Inherited;

  FreeAndNil(IgrUsbSerialList_);

  for i := 0 to DriveList_.Count - 1 do
  begin
//    if DriveList_[i].bReadOnly then
      SetReadOnly(DriveList_[i].DriveInfo.sDrive, DriveList_[i].DriveInfo.nDiskNum, false);
  end;

  FreeAndNil(qArrivalUsb_);
  if UsbNoti_ <> nil then
    FreeAndNil(UsbNoti_);
  FreeAndNil(EmDriveList_);
  FreeAndNil(DriveList_);
end;

procedure TThdUsbMon.OnUsbNotify(Sender: TObject; const Item: PUsbEnt; Action: TCollectionNotification);
begin
  if Action = cnRemoved then
    Dispose(Item);
end;

procedure TThdUsbMon.AddDrive(sDrive: String);
var
  i: Integer;
  pEnt: PUsbEnt;
begin
  try
    if GetDriveSize(sDrive) <> 0 then
    begin
      case GetDriveType(PChar(sDrive)) of
        DRIVE_FIXED, // 고정 디스크는 제외 25_0826 17:28:09 kku // Fixed 도 적용 되도록 롤백 (4005) 25_0917 17:39:41 kku
        DRIVE_REMOVABLE :
        begin
          pEnt := GetDriveEntByDrive(sDrive);
          if pEnt <> nil then
          begin
            Lock;
            try
              i := DriveList_.IndexOf(pEnt);
              if i <> -1 then
                DriveList_.Delete(i);
            finally
              Unlock;
            end;
          end;
          qArrivalUsb_.Enqueue(sDrive);

//          _Trace('UsbMon .. AddDrive() .. Drive=%s', [sDrive], 1);
        end;
      end;
    end;
  except
    on E: Exception do
      ETgException.TraceException(Self, E, 'Fail .. AddDrive()');
  end;
end;

procedure TThdUsbMon.AddEmptyDrive(sDrive: String);
var
  i: Integer;
  pEnt: PUsbEnt;
begin
  try
    if GetDriveSize(sDrive) = 0 then
    begin
      case GetDriveType(PChar(sDrive)) of
        DRIVE_FIXED,
        DRIVE_REMOVABLE :
        begin
          pEnt := GetDriveEntByDrive(sDrive);
          if pEnt <> nil then
          begin
            Lock;
            try
              i := DriveList_.IndexOf(pEnt);
              if i <> -1 then
                DriveList_.Delete(i);
            finally
              Unlock;
            end;
          end;

          if EmDriveList_.IndexOf(sDrive) = -1 then
            EmDriveList_.Add(sDrive);
        end;
      end;
    end;
  except
    on E: Exception do
      ETgException.TraceException(Self, E, 'Fail .. AddEmptyDrive()');
  end;
end;

procedure TThdUsbMon.DelDrive(sDrive: String);
var
  i: Integer;
  pEnt: PUsbEnt;
begin
  try
    pEnt := GetDriveEntByDrive(sDrive);
    if pEnt <> nil then
    begin
      if hRcvWnd_ <> 0 then
        SendMessage(hRcvWnd_, WM_USBCONTROL_NOTIFY, ACTION_USBCONTROL_REMOVE, NativeInt(pEnt));

      Lock;
      try
        i := DriveList_.IndexOf(pEnt);
        if i <> -1 then
        begin
        // QueryRemove에서 아래 처럼해도 안된다 22_0823 12:50:23 kku
    //      if ActionKind_ = upkReadOnly then
    //        SetReadOnly(pEnt.DriveInfo.sDrive, pEnt.DriveInfo.nDiskNum, false);

          DriveList_.Delete(i);
//          _Trace('UsbMon .. DelDrive() .. Drive=%s', [sDrive], 1);
        end;
      finally
        Unlock;
      end;
    end;

    i := EmDriveList_.IndexOf(sDrive);
    if (i > -1) and not DirectoryExists(sDrive) then
      EmDriveList_.Delete(i);
  except
    on E: Exception do
      ETgException.TraceException(Self, E, 'Fail .. DelDrive()');
  end;
end;

procedure TThdUsbMon.OnUSBArrival(Sender: TObject; pInfo: PDevBroadcastVolume);
var
  sDrive: String;
begin
  sDrive := GetDriveFromMask(pInfo.dwUnitmask);
//  if GetDriveExtent(sDrive).liExtentLength.QuadPart = 0 then
  if not DirectoryExists(sDrive) then
    AddEmptyDrive(sDrive)
  else
    AddDrive(sDrive);
end;

procedure TThdUsbMon.OnUSBQueryRemove(Sender: TObject; sDrive: String; var bAccept: Boolean);
begin
  bAccept := true;
end;

procedure TThdUsbMon.OnUSBRemove(Sender: TObject; pInfo: PDevBroadcastVolume);
var
  sDrive: String;
begin
  sDrive := GetDriveFromMask(pInfo.dwUnitmask);
  DelDrive(sDrive);
end;

function TThdUsbMon.GetDriveEntByDrive(sDrive: String): PUsbEnt;
var
  i: Integer;
begin
  Result := nil;
  Lock;
  try
    for i := 0 to DriveList_.Count - 1 do
      if CompareText(DriveList_[i].DriveInfo.sDrive, sDrive) = 0 then
        Exit(DriveList_[i]);
  finally
    Unlock;
  end;
end;

procedure TThdUsbMon.ProcessUsbControl(pEnt: PUsbEnt);
var
  sTemp: String;
  bIsExeDiskPart: Boolean;
  i: Integer;
begin
  try
    bIsExeDiskPart := false;

    case ActionKind_ of
      upkMonitor  :
      begin
  //      if hRcvWnd_ <> 0 then
  //        SendMessage(hRcvWnd_, WM_USBCONTROL_NOTIFY, ACTION_USBCONTROL_MONITOR, NativeInt(pEnt));
      end;
      upkReadOnly :
      begin
        if pEnt.DriveInfo.nDiskNum > 0 then
        begin
        {$IFDEF _HE_}
          var sVenders := gMgSvc.ModePolicy.UsbExceptVender;
          if (sVenders <> '') and (pEnt.DriveInfo.sFriendlyName <> '') then
          begin
            var VenderList: TStringList;
            Guard(VenderList, TStringList.Create);
            SplitString(UpperCase(sVenders), ';', VenderList);
            sTemp := UpperCase(pEnt.DriveInfo.sFriendlyName);
            for i := 0 to VenderList.Count - 1 do
              if sTemp.Contains(VenderList[i]) then
              begin
                if pEnt.bReadOnly and GetRecoverWB then
                begin
                  SetReadOnly(pEnt.DriveInfo.sDrive, pEnt.DriveInfo.nDiskNum, false);
                  pEnt.bReadOnly := false;
                end;
                exit;
              end;
          end;

          if sIgrUsbSerials_ <> gMgSvc.ModePolicy.IgrUsbSerials then
          begin
            sIgrUsbSerials_ := gMgSvc.ModePolicy.IgrUsbSerials;
            SplitString(UpperCase(sIgrUsbSerials_), ';', IgrUsbSerialList_);
          end;

          if IgrUsbSerialList_.Count > 0 then
          begin
            for i := 0 to IgrUsbSerialList_.Count - 1 do
            begin
              if Pos(IgrUsbSerialList_[i], UpperCase(pEnt.DriveInfo.sSerial)) > 0 then
              begin
                if pEnt.bReadOnly and GetRecoverWB then
                begin
                  SetReadOnly(pEnt.DriveInfo.sDrive, pEnt.DriveInfo.nDiskNum, false);
                  pEnt.bReadOnly := false;
                end;
                exit;
              end;
            end;
          end;
        {$ENDIF}

        // 쓰기 체크, 복구 기능 비활성 23_0324 11:56:03 kku
        {$IF true}
//          if ( {$IFDEF _HE_} (CUSTOMER_TYPE = CUSTOMER_GEC) or {$ENDIF}
//             (GetProcessPidByName('diskpart.exe') <> 0) ) and
//             ( (GetTickCount - pEnt.dwRoTick) > 3000 ) then

          bIsExeDiskPart := GetProcessPidByName('diskpart.exe') <> 0;
          if not pEnt.bReadOnly or bIsExeDiskPart or
             ( (GetTickCount - pEnt.dwRoTick) < 30000 ) then
          begin
            if bIsExeDiskPart then
              pEnt.dwRoTick := GetTickCount;

            if GetRecoverWB and (IsReadOnlyByWriteProbe(pEnt.DriveInfo.sDrive) <> 1) then
            begin
              SetReadOnly(pEnt.DriveInfo.sDrive, pEnt.DriveInfo.nDiskNum, true);
//              _Trace('UsbMon .. SetRO .. Drive = %s', [pEnt.DriveInfo.sDrive], 1);
              Sleep(1000);
              pEnt.bReadOnly := IsReadOnlyByWriteProbe(pEnt.DriveInfo.sDrive) = 1;
//              _Trace('UsbMon .. SetRO .. Drive = %s .. Result = %s', [pEnt.DriveInfo.sDrive, BooleanToStr(pEnt.bReadOnly, 'true', 'false')], 1);
            end else
            if not pEnt.bReadOnly then
              pEnt.bReadOnly := true;
          end;

          if not pEnt.bRoPopup then
          begin
        {$ELSE}
          if SetReadOnly and not pEnt.bRoPopup then
          begin
            SetReadOnly(pEnt.DriveInfo.sDrive, pEnt.DriveInfo.nDiskNum, true);
        {$IFEND}
            pEnt.bRoPopup := true;

          {$IFDEF _HE_}
            if gMgSvc.ModePolicy.UsbBlockKind = ubkReadOnly then
            begin
              with pEnt.DriveInfo do
              begin
                var sType: String := '';
                case GetDriveType(PChar(sDrive)) of
                  DRIVE_FIXED     : sType := 'FIXED';
                  DRIVE_CDROM     : sType := 'CDROM';
                  DRIVE_REMOVABLE : sType := 'REMOVABLE';
                  else sType := 'Unknown';
                end;

                if gMgSvc.ModePolicy.USBPopup then
                  gMgSvc.PopupMessage(TYPE_MSG_PREVENT_USBDISCONN, sDrive + '|' +
                      ByteSizeToStr(llSize) + '|' + sFriendlyName + '|' + sType + '|' + sSerial);

                var sMsg: String := Format('USB ReadOnly : Name=%s, Drive=%s (%s), Type=%s, Serial=%s',
                                      [sFriendlyName, sDrive, ByteSizeToStr(llSize), sType, sSerial]);
                if gMgSvc.IsNewApi then
                begin
                  var LogInfo: TLogInfo;
                  ZeroMemory(@LogInfo, SizeOf(LogInfo));
                  LogInfo.sCode := LOGCODE_PREVENT_USB;
                  LogInfo.sDevName := sFriendlyName;
                  LogInfo.sDevSerial := sSerial;
                  LogInfo.sDevClassId := sClassGuid;
                  LogInfo.sSummary := sMsg;
                  gMgSvc.SendEventLogEx(@LogInfo);
                end else
                  gMgSvc.SendEventLog(URI_USER_ACTION, LOGCODE_PREVENT_USB, sMsg);
              end;
            end;
          {$ENDIF}

            if hRcvWnd_ <> 0 then
              SendMessage(hRcvWnd_, WM_USBCONTROL_NOTIFY, ACTION_USBCONTROL_READONLY, NativeInt(pEnt));
          end;
        end;
      end;
      upkBlock    :
      begin
  //      if pEnt.UsbInfo.sSerial <> '' then // Fixed는 여기가 비어 있음 22_0812 13:51:40 kku
        begin
          var sRemoveSearial := EjectDrive(pEnt.DriveInfo.sDrive);
          if (sRemoveSearial <> '') and (sRemoveSearial <> FAIL_EJECT) then
          begin
  //          pEnt.UsbInfo.sSerial := sRemoveSearial;
            if hRcvWnd_ <> 0 then
              SendMessage(hRcvWnd_, WM_USBCONTROL_NOTIFY, ACTION_USBCONTROL_BLOCK, NativeInt(pEnt));
          end;
        end;
      end;
    end;
  except
    on E: Exception do
      ETgException.TraceException(Self, E, 'Fail .. ProcessUsbControl()');
  end;
end;

procedure TThdUsbMon.SetRecoverWB(bVal: Boolean);
begin
  Lock;
  try
    if bval <> bRecoverWB_ then
      bRecoverWB_ := bVal;
  finally
    Unlock;
  end;
end;

function TThdUsbMon.GetRecoverWB: Boolean;
begin
  Lock;
  try
    Result := bRecoverWB_;
  finally
    Unlock;
  end;
end;

procedure TThdUsbMon.UpdateDriveInfo;
var
  i: Integer;
  dwLogicalDrv: DWORD;
  sDrive: String;
  pEnt: PUsbEnt;
  pRec: PUSBRec;
  DriveInfo: TDriveInfo;
begin
  try
    dwLogicalDrv := GetLogicalDrives;
    for i := 0 to 31 do
      if (dwLogicalDrv and (1 shl i)) > 0 then
      begin
        sDrive := Format('%s:\', [Char(Integer('A')+i)]);
//        if GetDriveExtent(sDrive).liExtentLength.QuadPart = 0 then
        if not DirectoryExists(sDrive) then
        begin
          if EmDriveList_.IndexOf(sDrive) = -1 then
            EmDriveList_.Add(sDrive);
          continue;
        end;

        if UpCase(sDrive[1]) = 'C' then
          continue;

        if GetDriveType(PChar(sDrive)) <> DRIVE_REMOVABLE then
          continue;

        GetDriveDetail(sDrive, @DriveInfo);
      // USB인식 추가 보완
        if Pos('USB', UpperCase(DriveInfo.sSerial)) <> 1 then
          continue;

        New(pEnt);
        ZeroMemory(pEnt, SizeOf(TUsbEnt));
        pEnt.DriveInfo := DriveInfo;
        DriveList_.Add(pEnt);

        ProcessUsbControl(pEnt);
      end;
  except
    on E: Exception do
      ETgException.TraceException(Self, E, 'Fail .. UpdateDriveInfo()');
  end;
end;

procedure TThdUsbMon.Execute;
var
  pEnt: PUsbEnt;
  sDrive: String;
  i: Integer;
  IgrList: TStringList;
  bEmptyDrvArrival: Boolean;
begin
  UpdateDriveInfo;
  while not Terminated and not GetWorkStop do
  begin
    try
      Sleep(500);

      if qArrivalUsb_.Count = 0 then
      begin
        bEmptyDrvArrival := false;
        try
          if DriveList_ <> nil then
          begin
            for i := DriveList_.Count - 1 downto 0 do
            begin
              if not DirectoryExists(DriveList_[i].DriveInfo.sDrive) then
              begin
                case GetDriveType(PChar(sDrive)) of
                  DRIVE_FIXED,
                  DRIVE_REMOVABLE :
                  begin
                  // 빈드라이브 감지 24_0704 14:55:08 kku
                    if EmDriveList_.IndexOf(sDrive) = -1 then
                      EmDriveList_.Add(sDrive);

                  // 드라이브 해제 알림
                    if hRcvWnd_ <> 0 then
                      SendMessage(hRcvWnd_, WM_USBCONTROL_NOTIFY, ACTION_USBCONTROL_REMOVE, NativeInt(DriveList_[i]));
                  end;
                end;

                DriveList_.Delete(i);
                continue;
              end;

              ProcessUsbControl(DriveList_[i]);
            end;
          end;

          if EmDriveList_ <> nil then
          begin
            for i := EmDriveList_.Count -1 downto 0 do
            begin
              sDrive := EmDriveList_[i];
              if DirectoryExists(sDrive) and
                 (GetDriveEntByDrive(sDrive) = nil) then
              begin
                bEmptyDrvArrival := true;
                break;
              end;
            end;
          end;
        except
        // ...
        end;

        if not bEmptyDrvArrival then
          continue;
      end else
        sDrive := qArrivalUsb_.Dequeue;

      if not DirectoryExists(sDrive) then
        continue;

//      _Trace('UsbMon .. 1 .. Arrival Drive=%s', [sDrive], 1);

      New(pEnt);
      ZeroMemory(pEnt, SizeOf(TUsbEnt));
      pEnt.DriveInfo.sDrive := sDrive;
      GetDriveDetail(sDrive, @pEnt.DriveInfo);

      Lock;
      try
        DriveList_.Add(pEnt);
      finally
        Unlock;
      end;

    {$IFDEF _HE_}
      try
      // 읽기만 정책 사용중인지 상태체크
        pEnt.bReadOnly := IsReadOnlyByWriteProbe(pEnt.DriveInfo.sDrive) = 1;

//        _Trace('UsbMon .. 2 .. ReadOnly = %s', [BooleanToStr(pEnt.bReadOnly, 'true', 'false')], 1);

        if sIgrUsbSerials_ <> gMgSvc.ModePolicy.IgrUsbSerials then
        begin
          sIgrUsbSerials_ := gMgSvc.ModePolicy.IgrUsbSerials;
          SplitString(UpperCase(sIgrUsbSerials_), ';', IgrUsbSerialList_);
        end;

      // 읽기만 정책도 예외 추가 23_1218 13:59:17 kku
        if IgrUsbSerialList_.Count > 0 then
        begin
          var bIgrUsb: Boolean := false;
          for i := 0 to IgrUsbSerialList_.Count - 1 do
          begin
            if Pos(IgrUsbSerialList_[i], UpperCase(pEnt.DriveInfo.sSerial)) > 0 then
            begin
              bIgrUsb := true;
              break;
            end;
          end;

          if bIgrUsb then
            continue;
        end;
      except
      // ..
      end;
    {$ENDIF}

      if hRcvWnd_ <> 0 then
        SendMessage(hRcvWnd_, WM_USBCONTROL_NOTIFY, ACTION_USBCONTROL_ARRIVAL, NativeInt(pEnt));

      ProcessUsbControl(pEnt);
    except
      on E: Exception do
      begin
        ETgException.TraceException(Self, E, 'File .. Execute()');
        Sleep(1000);
      end;
    end;
  end;
end;

end.