unit DUsbMonMain;

interface

uses
  Winapi.Windows, Winapi.Messages, System.SysUtils, System.Variants, System.Classes, Vcl.Graphics,
  Vcl.Controls, Vcl.Forms, Vcl.Dialogs, Vcl.StdCtrls, Vcl.ExtCtrls, ThdUsbMon, Bs1FltCtrl,
  Tocsg.Trace;

type
  TDlgUsbMonMain = class(TForm)
    pnTop: TPanel;
    btnPreventUsb: TButton;
    rdReadOnly: TRadioButton;
    rdBlock: TRadioButton;
    mmLog: TMemo;
    rdMonitor: TRadioButton;
    edDrive: TEdit;
    btnDismount: TButton;
    btnDrvInfo: TButton;
    chkKernel: TCheckBox;
    rbUsbEnable: TRadioButton;
    rbUsbDisable: TRadioButton;
    btnBs1SetPolicy: TButton;
    grpBs1CtrlUsb: TGroupBox;

    procedure btnPreventUsbClick(Sender: TObject);
    procedure btnDismountClick(Sender: TObject);
    procedure btnDrvInfoClick(Sender: TObject);
    procedure btnBs1SetPolicyClick(Sender: TObject);
  private
    { Private declarations }
    ThdUsbMon_: TThdUsbMon;
    Trace_: TTgTrace;
  public
    { Public declarations }
    Bs1Ctrl: TBs1fltControl;
    kernelUsed_: Boolean;
    Constructor Create(aOwner: TComponent); override;
    Destructor Destroy; override;

    procedure process_WM_USBCONTROL_NOTIFY(var msg: TMessage); Message WM_USBCONTROL_NOTIFY;
  end;

var
  DlgUsbMonMain: TDlgUsbMonMain;

implementation

uses
  Tocsg.Driver, Tocsg.Convert, Tocsg.Path, Tocsg.Disk;

{$R *.dfm}


function MyLogCallback(Context: Pointer): DWORD; stdcall;
var
  logMsg: string;

begin
  // 로그 처리 로직 (Context는 문자열 포인터일 가능성이 높음)
  if Context = nil then Exit;

  LogMsg := string(PWideChar(Context));

  OutputDebugString(PWideChar(Context));

  TThread.Queue(nil,
    procedure
    begin
      // 폼이 살아있는지 확인 후 로그 추가
      if Assigned(DlgUsbMonMain) and Assigned(DlgUsbMonMain.mmLog) then
      begin
        if DlgUsbMonMain.mmLog.Lines.Count > 1000 then
          DlgUsbMonMain.mmLog.Lines.Delete(0);

        DlgUsbMonMain.mmLog.Lines.Add(LogMsg);
      end;
    end);

  Result := 0;
end;

Constructor TDlgUsbMonMain.Create(aOwner: TComponent);
var
  state: DWORD;
  path: string;
begin
  Inherited Create(aOwner);
  Trace_ := TTgTrace.Create(GetRunExePathDir, CutFileExt(GetRunExeName) + '.log');
  Trace_.Level := 10;
  ThdUsbMon_ := nil;
  kernelUsed_:= FALSE;
  rdBlock.Checked:= True;
  chkKernel.Checked:= True;

  path:= GetRunExePathDir;
  Bs1Ctrl := TBs1fltControl.Create;

  try
    // 커널 드라이버의 위치를 정한다.
    state := Bs1Ctrl.Init(path, MyLogCallback);

    if state = 0 then
    begin
//      ShowMessage('Bs1Ctrl 초기화 성공');
      mmLog.Lines.Add('Bs1Ctrl 초기화 성공');
      Bs1Ctrl.Bs1FltBegin(1);
    end
    else
    begin
      mmLog.Lines.Add('Bs1Ctrl 초기화 실패: ' + IntToStr(state));
    end;

  finally
    // 전역 변수로 관리하거나 필요 없으면 해제
    // Bs1Ctrl.Free;
  end;

end;

Destructor TDlgUsbMonMain.Destroy;
begin
  if ThdUsbMon_ <> nil then
    FreeAndNil(ThdUsbMon_);

  if Bs1Ctrl <> nil then
  begin
    if Assigned(Bs1Ctrl.Bs1FltCleanup) then
      Bs1Ctrl.Bs1FltCleanup;

    Bs1Ctrl.Free;
  end;

  Inherited;
  FreeAndNil(Trace_);
end;

procedure TDlgUsbMonMain.btnBs1SetPolicyClick(Sender: TObject);
var
  state: enum_devicestate;
  protect: DWORD;
begin
  if rbUsbEnable.Checked then
  begin
    state:= dsEnable;
    protect:= 0;
  end
  else
  begin
    state:= dsDisable;
    protect:= 1;
  end;

  if not Assigned(Bs1Ctrl.Bs1FltCleanup) then
    exit;

  Bs1Ctrl.Bs1fltSetDeviceProtect(protect);
  Bs1Ctrl.Bs1FltSetPolicy(DWORD(BDC_USB_DISK), DWORD(state), 1);
  Bs1Ctrl.Bs1FltSetPolicy(DWORD(BDC_EXTERNALHDD), DWORD(state), 1);

end;

procedure TDlgUsbMonMain.btnDismountClick(Sender: TObject);
begin
  edDrive.Text := Trim(edDrive.Text);
  if (edDrive.Text = '') and (not DirectoryExists(edDrive.Text)) then
  begin
    ShowMessage('실패 1');
    exit;
  end;

//  ShowMessage(GetVolumeName(edDrive.Text));
//  exit;

  if EjectDrive(edDrive.Text, nil, true, true) = FAIL_EJECT then
    ShowMessage('실패 2')
  else
    ShowMessage('성공');
end;

procedure TDlgUsbMonMain.btnDrvInfoClick(Sender: TObject);
var
  DriveInfo: TDriveInfo;
begin
  edDrive.Text := Trim(edDrive.Text);
  if (edDrive.Text = '') and (not DirectoryExists(edDrive.Text)) then
  begin
    ShowMessage('실패 1');
    exit;
  end;


  GetDriveDetail(edDrive.Text, @DriveInfo, true);
  mmLog.Lines.Add('Drive : ' + DriveInfo.sDrive);
  mmLog.Lines.Add('Drive Type : ' + GetDriveTypeToStr(GetDriveType(PChar(DriveInfo.sDrive))));
  mmLog.Lines.Add('Serial : ' + DriveInfo.sSerial);
  mmLog.Lines.Add('Class : ' + DriveInfo.sClass);
  mmLog.Lines.Add('ClassGuid : ' + DriveInfo.sClassGuid);
  mmLog.Lines.Add('Description : ' + DriveInfo.sDesc);
  mmLog.Lines.Add('FriendlyName : ' + DriveInfo.sFriendlyName);
  mmLog.Lines.Add(Format('Size : %s (%d)', [ByteSizeToStr(DriveInfo.llSize), DriveInfo.llSize]));
  mmLog.Lines.Add('DiskNum : ' + IntToStr(DriveInfo.nDiskNum));
end;

procedure TDlgUsbMonMain.btnPreventUsbClick(Sender: TObject);
var
  ActionKind: TUsbProcKind;
begin
  if ThdUsbMon_ = nil then
  begin
    if rdMonitor.Checked then
      ActionKind := upkMonitor
    else if rdReadOnly.Checked then
      ActionKind := upkReadOnly
    else
      ActionKind := upkBlock;

    mmLog.Clear;
    ThdUsbMon_ := TThdUsbMon.Create(Handle, ActionKind);
    ThdUsbMon_.StartThread;
    if chkKernel.Checked then
    begin
      kernelUsed_:= True;
      Bs1Ctrl.Bs1fltSetDeviceProtect(1);
      Bs1Ctrl.Bs1FltSetPolicy(DWORD(BDC_USB_DISK), DWORD(dsDisable), 1);
      Bs1Ctrl.Bs1FltSetPolicy(DWORD(BDC_EXTERNALHDD), DWORD(dsDisable), 1);
    end
    else
     kernelUsed_:= False;

     chkKernel.Enabled:= False;


  end else begin
    if MessageBox(Handle, PChar('중지하시겠습니까?'),
      PChar(Caption), MB_ICONQUESTION or MB_YESNO) = IDNO then exit;

    FreeAndNil(ThdUsbMon_);
    if kernelUsed_ then
    begin
      Bs1Ctrl.Bs1fltSetDeviceProtect(0);
      Bs1Ctrl.Bs1FltSetPolicy(DWORD(BDC_USB_DISK), DWORD(dsEnable), 0);
      Bs1Ctrl.Bs1FltSetPolicy(DWORD(BDC_EXTERNALHDD), DWORD(dsEnable), 0);
    end;

    chkKernel.Enabled:= True;
  end;

  rdMonitor.Enabled := ThdUsbMon_ = nil;
  rdReadOnly.Enabled := rdMonitor.Enabled;
  rdBlock.Enabled := rdMonitor.Enabled;
  if rdMonitor.Enabled then
    btnPreventUsb.Caption := 'USB 제어 시작'
  else
    btnPreventUsb.Caption := 'USB 제어 중지';

  Application.ProcessMessages;
end;

procedure TDlgUsbMonMain.process_WM_USBCONTROL_NOTIFY(var msg: TMessage);

  procedure WriteLog(sLog: String);
  begin
    mmLog.Lines.Add(Format('[%s] %s', [DateTimeToStr(Now), sLog]));
  end;

var
  pEnt: PUsbEnt;
begin
  pEnt := PUsbEnt(msg.LParam);
  case msg.WParam of
    ACTION_USBCONTROL_ARRIVAL         :
    begin
      WriteLog(Format('연결됨 : Drive=%s, Size=%s, DevName=%s',
        [pEnt.DriveInfo.sDrive, ByteSizeToStr(pEnt.DriveInfo.llSize), pEnt.DriveInfo.sFriendlyName]));
    end;
    ACTION_USBCONTROL_REMOVE          :
    begin
      WriteLog(Format('해제됨 : Drive=%s, Size=%s, DevName=%s',
        [pEnt.DriveInfo.sDrive, ByteSizeToStr(pEnt.DriveInfo.llSize), pEnt.DriveInfo.sFriendlyName]));
    end;
    ACTION_USBCONTROL_READONLY        :
    begin
      WriteLog(Format('읽기전용 적용! : Drive=%s, Size=%s, DevName=%s',
        [pEnt.DriveInfo.sDrive, ByteSizeToStr(pEnt.DriveInfo.llSize), pEnt.DriveInfo.sFriendlyName]));
    end;
    ACTION_USBCONTROL_READONLY_RETRAY : ;
    ACTION_USBCONTROL_BLOCK           :
    begin
      WriteLog(Format('차단!! : Drive=%s, Size=%s, DevName=%s',
        [pEnt.DriveInfo.sDrive, ByteSizeToStr(pEnt.DriveInfo.llSize), pEnt.DriveInfo.sFriendlyName]));
    end;
  end;
end;

end.