Flash Drive Locking Script

[Solarlight27]

Ever wanted to use Autoit to lock up your flash drive at home?

Well here's the code for it. Code from user LvlUp and some minor modification from me.

All props to LvlUp and his coding abilities and assistance with this code.

Here's the main code:

expandcollapse popup

#NoTrayIcon
#include <Misc.au3>
#include <eject.au3>

Opt("TrayMenuMode",1)
$menu = TrayCreateItem("DriveLocker")
$menu2 = TrayCreateItem("")
$menuLock = TrayCreateItem("Lock Drive")
$menuExit = TrayCreateItem("Exit")
TraySetState()

While 1
$letter = DriveGetDrive("REMOVABLE")
If IsArray($letter) Then
  For $i = 1 to $letter[0]
   _Check($letter[$i])
  Next
EndIf
$time = TimerInit()
While TimerDiff($time) < 10000
  $msg = TrayGetMsg()
  TrayItemSetState($msg,4)
  Select
   Case $msg = $menuLock
   _LockDrive()
   Case $msg = $menuExit
   Exit
  EndSelect
  Sleep(100)
WEnd
WEnd

Func _Check($strDriveLetter)
If FileExists($strDriveLetter & "\IAmLocked") Then _Unlock($strDriveLetter)
EndFunc

Func _Unlock($strDriveLetter)
$strPassword = InputBox("Security Check", "Enter your password.", "", "*", "", "", @DesktopWidth/2, @DesktopHeight/2,15, "passwordBox")
ToolTip("Enter Password NOW.")
_MouseTrap(@DesktopWidth, @DesktopHeight-100)
If ($strPassword = "password") Then
  _MouseTrap()
  FileDelete($strDriveLetter & "\IAmLocked")
  _ReAnimate($strDriveLetter)
Else
  EjectVolume($strDriveLetter)
  _MouseTrap()
EndIf
EndFunc

Func _ReAnimate($strDriveLetter)
;Closes the autoplay window if open
WinClose("AutoPlay")
EndFunc

Func _LockDrive()
$letter = DriveGetDrive("REMOVABLE")
If IsArray($letter) Then
  For $i = 1 To $letter[0]
   If MsgBox(36,"Lock This Drive?","Do you wish to lock the " & StringUpper($letter[$i]) & " drive?") = 6 Then
    FileWrite($letter[$i] & "\IAmLocked","Locked")
    EjectVolume($letter[$i])
   EndIf
  Next
EndIf
EndFunc

And the includes file for eject.au3

expandcollapse popup

#Region ;**** Directives created by AutoIt3Wrapper_GUI ****
#AutoIt3Wrapper_Res_requestedExecutionLevel=asInvoker
#EndRegion ;**** Directives created by AutoIt3Wrapper_GUI ****
#include <WinApi.au3>
;Prototypes
;BOOL EjectVolume(TCHAR cDriveLetter);
;HANDLE OpenVolume(TCHAR cDriveLetter);
;BOOL LockVolume(HANDLE hVolume);
;BOOL DismountVolume(HANDLE hVolume);
;BOOL PreventRemovalOfVolume(HANDLE hVolume, BOOL fPrevent);
;BOOL AutoEjectVolume(HANDLE hVolume);
;BOOL CloseVolume(HANDLE hVolume);
;StringFormat Output
$szVolumeFormat = "\\\\.\\%s"
$szRootFormat = "%s\\"
$szErrorFormat = "Error %d: %s\n"
;------------------------------------------
;Arbitrary variables
;;Global Const $INVALID_HANDLE_VALUE = 0
;------------------------------------------
;DeviceIoControl Contants
Global Const $FSCTL_LOCK_VOLUME = int(0x090018)
Global Const $FSCTL_DISMOUNT_VOLUME = int(0x00090020)
Global Const $IOCTL_STORAGE_EJECT_MEDIA  = int(0x002D4808)
Global Const $IOCTL_STORAGE_MEDIA_REMOVAL = int(0x002D4804)
;------------------------------------------
;Retry Constants
Global Const $LOCK_TIMEOUT = 10000   ; 10 Seconds
Global Const $LOCK_RETRIES = 20
;$OpenVolume = InputBox("Ejecting...", "Enter the drive to eject", "G:")
;ConsoleWrite("Trying to Eject the drive " & EjectVolume($OpenVolume) & @crlf)
Func ReportError($szMsg)
    ConsoleWrite(StringFormat($szErrorFormat, _WinAPI_GetLastErrorMessage (), $szMsg) & @CRLF)
    Exit
EndFunc   ;==>ReportError
Func OpenVolume($cDriveLetter)
    ;HANDLE hVolume
    ;UINT uDriveType
    ;TCHAR szVolumeName[8]
    ;TCHAR szRootName[5]
    ;DWORD dwAccessFlags
    $szRootName = StringFormat($szRootFormat, $cDriveLetter)
    $uDriveType = DriveGetType($szRootName);
    ConsoleWrite($szRootName & @tab & $uDriveType & @crlf)
    Switch $uDriveType
        Case "Removable"
            $dwAccessFlags = 6
        Case "CDROM"
            $dwAccessFlags = 2
        Case Else
            ConsoleWrite("Cannot eject.  Drive type is incorrect." & @CRLF)
            Return $INVALID_HANDLE_VALUE
    EndSwitch
    $szVolumeName = StringFormat($szVolumeFormat, $cDriveLetter)
    ;$szVolumeName = $szVolumeFormat & $cDriveLetter
    ConsoleWrite($szVolumeName & @crlf )
    $hVolume = _WinAPI_CreateFile ($szVolumeName, 2,$dwAccessFlags, 6)
    #cs
        hVolume = CreateFile(   szVolumeName,
        dwAccessFlags,
        FILE_SHARE_READ | FILE_SHARE_WRITE,
        NULL,
        OPEN_EXISTING,
        0,
        NULL );
    #ce
    If ($hVolume == $INVALID_HANDLE_VALUE) Then ReportError("CreateFile");
    Return $hVolume;
EndFunc   ;==>OpenVolume
Func CloseVolume($hVolume)
    Return _WinAPI_CloseHandle ($hVolume);
EndFunc   ;==>CloseVolume
Func LockVolume($hVolume)
    Local $dwBytesReturned
    Local $dwSleepAmount
    Local $nTryCount
    local $iRead
    $dwSleepAmount = $LOCK_TIMEOUT / $LOCK_RETRIES;
    ; Do this in a loop until a timeout period has expired
    For $nTryCount = 0 To $nTryCount < $LOCK_RETRIES
        If _Device_Control($hVolume, $FSCTL_LOCK_VOLUME, $iRead) Then
            Return True
        Else
            Sleep($dwSleepAmount);
        EndIf
    Next
    Return False;
EndFunc   ;==>LockVolume
Func DismountVolume($hVolume)
    ConsoleWrite("Dismount " & $hVolume & @crlf)
    Local $dwBytesReturned, $iRead
    local $aResult = _Device_Control($hVolume, $FSCTL_DISMOUNT_VOLUME, $iRead)
    ;msgbox(0,"",$aResult)
    Return $aResult
    ;Return $dwBytesReturned
EndFunc   ;==>DismountVolume
Func PreventRemovalOfVolume($hVolume, $fPreventRemoval)
    Local $dwBytesReturned
    Local $aResult
    Local $lpInbuffer,$nInBufferSize,$lpOutBuffer,$nOutBufferSize,$lpOverlapped
    $PMRBUFFER = DllStructCreate("bool PreventMediaRemoval")
    DllStructSetData($PMRBUFFER,"PreventMediaRemoval",$fPreventRemoval)
    $lpBytesReturned    = DllStructCreate("int Read")
    $pRead   = DllStructGetPtr($lpBytesReturned, "Read")
    $aResult = Dllcall("kernel32.dll","int","DeviceIoControl","hwnd",$hVolume,"uint",$IOCTL_STORAGE_MEDIA_REMOVAL,"ptr",DllStructGetPtr($PMRBUFFER),"uint",DllStructGetSize($PMRBUFFER), _
    "ptr",$lpOutBuffer,"uint",$nOutBufferSize,"ptr",$pRead,"ptr",$lpOverlapped)
    if $aResult = 0 then msgbox(0,"",_WinAPI_GetLastErrorMessage())
    Return $aResult <> 0
     ;& PMRBuffer, sizeof (PREVENT_MEDIA_REMOVAL),
    ;NULL, 0,
    ; & dwBytesReturned,
    ;NULL);
EndFunc   ;==>PreventRemovalOfVolume
Func AutoEjectVolume($hVolume)
    Local $aResult, $iRead;
    $aResult = _Device_Control($hVolume, $IOCTL_STORAGE_EJECT_MEDIA, $iRead)
    Return $aResult
EndFunc   ;==>AutoEjectVolume
Func EjectVolume($cDriveLetter)
    Local $hVolume;
    Local $fRemoveSafely = False;
    Local $fAutoEject = False;
    ; Open the volume.
    $hVolume = OpenVolume($cDriveLetter);
    If $hVolume == $INVALID_HANDLE_VALUE Then Return False
    ; Lock and dismount the volume.
    If LockVolume($hVolume) And DismountVolume($hVolume) Then
        $fRemoveSafely = True;
        ConsoleWrite("Volume Locked and Dismounted, trying to Eject " & @crlf)
        ; Set prevent removal to false and Eject the volume.
        If PreventRemovalOfVolume($hVolume, False) And AutoEjectVolume($hVolume) Then
            $fAutoEject = True;
        EndIf
    Else
        ConsoleWrite("Volume can't be locked or dismounted, please close possible opened files" & @crlf)
    EndIf
    ; Close the volume so other processes can use the drive.
    If CloseVolume($hVolume) = False Then
        Return False;
    EndIf
    If $fAutoEject Then
        ConsoleWrite(StringFormat("Media in drive %s has been ejected safely.\n", $cDriveLetter))
    Else
        If $fRemoveSafely Then
            ConsoleWrite(StringFormat("Media in drive %s can be safely removed.\n", $cDriveLetter))
        EndIf
    EndIf
    Return True;
EndFunc   ;==>EjectVolume
Func _Device_Control($hDevice, $dwIoControlAutoit, ByRef $iRead)
    Local $aResult
    Local $lpInbuffer,$nInBufferSize,$lpOutBuffer,$nOutBufferSize,$lpOverlapped
    $tRead   = DllStructCreate("int Data")
    $aResult = Dllcall("kernel32.dll","int","DeviceIoControl","hwnd",$hDevice,"uint",$dwIoControlAutoit,"ptr",$lpInBuffer,"uint",0, _
    "ptr",$lpOutBuffer,"uint",0,"ptr",DllStructGetPtr($tRead),"ptr",$lpOverlapped)
    $iRead   = DllStructGetData($tRead, "Data")
    ConsoleWrite("Device Control " & $iRead & @CRLF)
    Return $aResult<>0
EndFunc   ;==>_Device_Control

Enjoy.

[llewxam]

I just discovered this post and LockVolume and DismountVolume look like exactly the thing I've been looking for! Thanks for posting!

Ian