AutoHotkey Community

/*
NOTE: You need to customize drives to ignore in PopUSB subroutine.

A lot of the code (the complicated parts) are based on code by SKAN including
his Safely Remove USB Flash Drive - 45L:
http://www.autohotkey.com/forum/viewtopic.php?t=44873

The script also heavily relies on TrayIcon.ahk by Sean:
http://www.autohotkey.com/forum/viewtopic.php?t=18652

And it uses a pipe script from HotkeyIt/Lexikos:
;Below from is from a post by HotkeyIt: http://www.autohotkey.com/forum/viewtopic.php?p=263671#263671
;He referenced Lexikos as the main source: http://www.autohotkey.com/forum/viewtopic.php?t=25867

To use the UnlockIt subroutine, Unlocker has to be installed:
http://ccollomb.free.fr/unlocker/
(this is a temporary solution - I want to ultimately use AHK to find the offending processes)

You may also want to customize the auto-pop-up behavior of the ejection box, which is loosely
based on bmcclure post here:
http://www.autohotkey.com/forum/viewtopic.php?p=226847#226847
*/
#NoEnv
#SingleInstance, Force
SetBatchLines, -1
;SetWinDelay -1
;SetKeyDelay -1
;SetControlDelay -1
DetectHiddenWindows, On
SetTitleMatchMode, 2
SendMode, Input
CoordMode, Mouse, Screen
SetWorkingDir % A_ScriptDir
Process,Priority,,B
OnExit,ExitCleanUp

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;   Start of USB Remove initiation
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;below list floppy drives, and fixed,virtual,ram hard drives to skip (all permanent drives, that may appear to AHK as Removable or Fixed e.g. RAM drives often emulate fixed drives)
SkipDrives = (A|B|C|Q|R|Z)

;record script's PID for use in auto-closing its pop-ups
pidParent := DllCall("GetCurrentProcessId")

;record system tray id for clicking Safely Remove Hardware
hSysTray := WinExist( "ahk_class SystemTray_Main" )

;set script to receive notification on mouse-over (tray icon)
OnMessage(0x404,"AutoHotkey_Notify")

;hide the default "Safely Remove Hardware" tray icon (uID = 1226)
If RegExMatch(TrayIcons("explorer.exe"), "(?<=idn: )(?P<Idn>\d+) \| Pid: (?P<Pid>\d+) \| uID: 1226", srh)
  HideTrayIcon(srhIdn)

;Context menus cannot be timed to self-terminate, so the auxiliary script below is launched via a pipe to wait for a context menu window from parent script's process (PID) right before parent shows menu, and to close it within 1.5s if/once not in use (mouse cursor not over it)
script =
(
  #NoTrayIcon
  DetectHiddenWindows, On
  WinWait, ahk_class #32768 ahk_pid %pidParent%,,1
  If hPop := WinExist("ahk_class #32768 ahk_pid %pidParent%")
  {
    Loop,4
    {
      Sleep, 1500   
      IfWinNotExist, ahk_id `%hPop`%
        ExitApp
      MouseGetPos,,,mPop,,1
      If (mPop != hPop)
        Break
    }
    IfWinExist, ahk_id `%hPop`%
      WinClose
  }
  ExitApp
)

Menu, Tray, UseErrorLevel
Menu, Tray, Icon, Hotplug.dll, 2
Menu, Tray, Add,
Menu, Tray, Add, Show Removable Devices, PopUSB
Menu, USB, UseErrorLevel
Return

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;   End of USB Remove initiation
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

#u::
;Specify drives to ignore below
PopUSB:
  DriveGet,RDRV,List,REMOVABLE
  RDRV := RegExReplace(RDRV,SkipDrives)
  DriveGet,FDRV,List,FIXED
  ADRV := RDRV . RegExReplace(FDRV,SkipDrives)
  ;if no ejectable drives, no menu
  If !ADRV
    Return

  ejectionAttempt := 0
  ;skip menu if only one drive to eject and ejection was hotkey initiated
  If (A_ThisHotkey == "#u" && A_TimeSinceThisHotkey < 100 && StrLen(ADRV) == 1)
    Gosub, EjectDirectly
  Else
  {
    Menu, USB, DeleteAll
    Loop, Parse, ADRV
    {
      DriveGet, Label, Label, %A_LoopField%:   
      DriveGet, Size, Capacity, %A_LoopField%:
      Capacity := DriveSpace( A_LoopField,2 ), VarSetCapacity( DiskSz,16,0 )
      DllCall( "shlwapi.dll\StrFormatByteSize64A", Int64,Capacity, Str,DiskSz, UInt,16 )
      Menu, USB, Add, &%A_LoopField%: %Label%`t%DiskSz%`t, EjectUSB
    }
    RunScript(script)
    Menu, USB, Show
  }
Return

EjectDirectly:
EjectUSB:
Loop, 3
{
  ejectionAttempt++
  If (ejectionAttempt == 1)
  {
    Drv := (A_ThisLabel == "EjectDirectly") ? ADRV : SubStr( A_ThisMenuItem,2,1 ), noSafelyRemove := 1, hProblem := ""
    ;If drive is "removable" (flash) use SKAN's USBD_SafelyRemove
    IfInString, RDRV, %Drv%
      noSafelyRemove := USBD_SafelyRemove( Drv . ":" )
  }
  Else Sleep, 3000 ;wait 3 second and try again
 
  ;If drive not "removable", or USBD_SafelyRemove failed, try to eject drive from Windows "Safely Remove Hardware" tray menu
  If noSafelyRemove
  {
    BlockInput, on
    MouseGetPos,oX,oY
    MouseMove,%A_ScreenWidth%,%A_ScreenHeight%,0
    PostMessage, 1226, 1226, 0x201,,ahk_id %hSysTray% ; Left Click down
    PostMessage, 1226, 1226, 0x202,,ahk_id %hSysTray% ; Left Click Up
    WinWaitActive, ahk_id %hSysTray%,,5                ; Wait for SRH Tray left-click-Menu
   
    If (ejectionAttempt == 1)
    {
      ; MN_GETHMENU : Code for retrieving popup menu text adapted from Sean's following post
      ;   Get Info from Context Menu: www.autohotkey.com/forum/viewtopic.php?p=137692#137692
      SendMessage,0x1E1,0,0,,ahk_class #32768
      hMenu := ErrorLevel, USBcount := DllCall( "GetMenuItemCount",UInt,hMenu )
      Loop, % USBcount
      {
        idx := A_Index-1, idn := DllCall( "GetMenuItemID", UInt,hMenu, Int,idx ), nSize := DllCall( "GetMenuString", UInt,hMenu, Int,idx, Int,0, Int,0, UInt,0x400 ) + 1, mStr := ""
        VarSetCapacity( mStr,nSize )
        DllCall( "GetMenuString", UInt,hMenu, Int,idx, Str,mStr, Int,nSize, UInt,0x400 )
        If InStr( mStr, Drv . ":")
        {
          steps := A_Index
          Break
        }
      }
    }
    ControlSend,,{Down %steps%}{Enter},ahk_id %hSysTray%
    MouseMove,%oX%,%oY%,0
    BlockInput, Off
    ;loop 10 sec waiting for either the ejected drive to disappear (success), or 'Problem Ejecting' window to appear
    Loop, 10
    {
      WinWait,Problem Ejecting,,1
      If (!Errorlevel || !FileExist(Drv . ":"))
        Break
    }
    If hProblem := WinExist("Problem Ejecting")
    {
      ControlClick, Button1, ahk_id %hProblem%
      If (ejectionAttempt < 3) ;try again
        Continue
      MsgBox,4,Device %Drv%: locked,Do you want to run Unlocker?`n(Requires Unlocker to be installed),11
      IfMsgBox Yes
        Run, %A_ProgramFiles%\Unlocker\Unlocker.exe %Drv%:
    }
    Else Break
  }
  Else Break
}
;cleanup: 
IfExist, %Drv%:
  TrayTip,Mystery error,Ejection of device %Drv%: failed.`n`nhProblem: %hProblem%`nejectionAttempt: %ejectionAttempt%`nnoSafelyRemove: %noSafelyRemove%,11
If noSafelyRemove
  SetTimer,HideIcon,-11000
Return

HideIcon:
RegExMatch(TrayIcons("explorer.exe"), "(?<=idn: )(?P<Idn>\d+) \| Pid: (?P<Pid>\d+) \| uID: 1226", srh)
  HideTrayIcon(srhIdn)
Return

;opens tray menu/s when mouse is over icon and interprets clicks
AutoHotkey_Notify(wParam, lParam)
{
  static ignoreNext
  If (lParam = 0x205) ;Right-click menu
  {
    ignoreNext := A_Now
    Menu, Tray, Show
  }
  Else If (lParam = 0x202) ;Left-click menu
  {
    ignoreNext := A_Now
    KeyWait, LButton, D T0.2  ; wait to see if 2nd left-click follows
    If ErrorLevel
      Gosub, PopUSB ;no second click, execute default left-click action
    Else ListVars ;double-click received
  }
  Else
  {
    ignoreWait := ignoreNext
    EnvSub,ignoreWait,A_Now,S
    If (ignoreWait < -4 || !ignoreWait)
    {
      ignoreNext := A_Now
      Gosub, PopUSB
    }
  }
  Return
}

;unhides Windows 'Safely Remove Hardware' tray icon
ExitCleanUp:
  RegExMatch(TrayIcons("explorer.exe"), "s)(?<=idn: )\d+(?= \| Pid: \d+ \| uID: 1226)", idn), HideTrayIcon(idn,False)
ExitApp
 
 
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;Functions below are 100% by other people and normally reside in my .\AutoHotkey\Lib direcorty:
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

; www.autohotkey.com/forum/viewtopic.php?p=92483#92483
DriveSpace(Drv="", Free=1)
{
  Drv.= ":\", VarSetCapacity(SPC, 30, 0), VarSetCapacity(BPS, 30, 0)
  VarSetCapacity(FC , 30, 0), VarSetCapacity(TC , 30, 0) 
  DllCall( "GetDiskFreeSpaceA",Str,Drv,UIntP,SPC,UIntP,BPS,UIntP,FC,UIntP,TC )
  Return Free=1 ? (SPC*BPS*FC) : (SPC*BPS*TC) ; Ternary Operator requires 1.0.46+
}

;from TrayIcon.ahk by Sean: http://www.autohotkey.com/forum/viewtopic.php?t=18652
TrayIcons(sExeName = "")
{
   WinGet,   pidTaskbar, PID, ahk_class Shell_TrayWnd
   hProc:=   DllCall("OpenProcess", "Uint", 0x38, "int", 0, "Uint", pidTaskbar)
   pProc:=   DllCall("VirtualAllocEx", "Uint", hProc, "Uint", 0, "Uint", 32, "Uint", 0x1000, "Uint", 0x4)
   idxTB:=   GetTrayBar()
      SendMessage, 0x418, 0, 0, ToolbarWindow32%idxTB%, ahk_class Shell_TrayWnd   ; TB_BUTTONCOUNT
   Loop,   %ErrorLevel%
   {
      SendMessage, 0x417, A_Index-1, pProc, ToolbarWindow32%idxTB%, ahk_class Shell_TrayWnd   ; TB_GETBUTTON
      VarSetCapacity(btn,32,0), VarSetCapacity(nfo,32,0)
      DllCall("ReadProcessMemory", "Uint", hProc, "Uint", pProc, "Uint", &btn, "Uint", 32, "Uint", 0)
         iBitmap   := NumGet(btn, 0)
         idn   := NumGet(btn, 4)
         Statyle := NumGet(btn, 8)
      If   dwData   := NumGet(btn,12)
         iString   := NumGet(btn,16)
      Else   dwData   := NumGet(btn,16,"int64"), iString:=NumGet(btn,24,"int64")
      DllCall("ReadProcessMemory", "Uint", hProc, "Uint", dwData, "Uint", &nfo, "Uint", 32, "Uint", 0)
      If   NumGet(btn,12)
         hWnd   := NumGet(nfo, 0)
      ,   uID   := NumGet(nfo, 4)
      ,   nMsg   := NumGet(nfo, 8)
      ,   hIcon   := NumGet(nfo,20)
      Else   hWnd   := NumGet(nfo, 0,"int64"), uID:=NumGet(nfo, 8), nMsg:=NumGet(nfo,12)
      WinGet, pid, PID,              ahk_id %hWnd%
      WinGet, sProcess, ProcessName, ahk_id %hWnd%
      WinGetClass, sClass,           ahk_id %hWnd%
      If !sExeName || (sExeName = sProcess) || (sExeName = pid)
         VarSetCapacity(sTooltip,128), VarSetCapacity(wTooltip,128*2)
      ,   DllCall("ReadProcessMemory", "Uint", hProc, "Uint", iString, "Uint", &wTooltip, "Uint", 128*2, "Uint", 0)
      ,   DllCall("WideCharToMultiByte", "Uint", 0, "Uint", 0, "str", wTooltip, "int", -1, "str", sTooltip, "int", 128, "Uint", 0, "Uint", 0)
      ,   sTrayIcons .= "idx: " . A_Index-1 . " | idn: " . idn . " | Pid: " . pid . " | uID: " . uID . " | MessageID: " . nMsg . " | hWnd: " . hWnd . " | Class: " . sClass . " | Process: " . sProcess . "`n" . "   | Tooltip: " . sTooltip . "`n"
   }
   DllCall("VirtualFreeEx", "Uint", hProc, "Uint", pProc, "Uint", 0, "Uint", 0x8000)
   DllCall("CloseHandle", "Uint", hProc)
   Return   sTrayIcons
}

HideTrayIcon(idn, bHide = True)
{
   idxTB := GetTrayBar()
   SendMessage, 0x404, idn, bHide, ToolbarWindow32%idxTB%, ahk_class Shell_TrayWnd   ; TB_HIDEBUTTON
   SendMessage, 0x1A, 0, 0, , ahk_class Shell_TrayWnd
}

GetTrayBar()
{
   ControlGet, hParent, hWnd,, TrayNotifyWnd1  , ahk_class Shell_TrayWnd
   ControlGet, hChild , hWnd,, ToolbarWindow321, ahk_id %hParent%
   Loop
   {
      ControlGet, hWnd, hWnd,, ToolbarWindow32%A_Index%, ahk_class Shell_TrayWnd
      If  Not   hWnd
         Break
      Else If   hWnd = %hChild%
      {
         idxTB := A_Index
         Break
      }
   }
   Return   idxTB
}

;Original script by Lexikos: Lexikos as the main source: http://www.autohotkey.com/forum/viewtopic.php?t=25867
;The function version below is from a post by HotkeyIt: http://www.autohotkey.com/forum/viewtopic.php?p=263671#263671
RunScript(TempScript, name="")
{
   pipe_name := name="" ? A_TickCount : name
   ; Before reading the file, AutoHotkey calls GetFileAttributes(). This causes
   ; the pipe to close, so we must create a second pipe for the actual file contents.
   ; Open them both before starting AutoHotkey, or the second attempt to open the
   ; "file" will be very likely to fail. The first created instance of the pipe
   ; seems to reliably be "opened" first. Otherwise, WriteFile would fail.
   pipe_ga := CreateNamedPipe(pipe_name, 2)
   pipe    := CreateNamedPipe(pipe_name, 2)
   if (pipe=-1 or pipe_ga=-1)
   {
      MsgBox CreateNamedPipe failed.
      ExitApp
   }
   Run, %A_AhkPath% "\\.\pipe\%pipe_name%",,,PID
   ; Wait for AutoHotkey to connect to pipe_ga via GetFileAttributes().
   DllCall("ConnectNamedPipe","uint",pipe_ga,"uint",0)
   ; This pipe is not needed, so close it now. (The pipe instance will not be fully
   ; destroyed until AutoHotkey also closes its handle.)
   DllCall("CloseHandle","uint",pipe_ga)
   ; Wait for AutoHotkey to connect to open the "file".
   DllCall("ConnectNamedPipe","uint",pipe,"uint",0)
   ; AutoHotkey reads the first 3 bytes to check for the UTF-8 BOM "". If it is
   ; NOT present, AutoHotkey then attempts to "rewind", thus breaking the pipe.
   Script:= chr(239) chr(187) chr(191) TempScript
   
   if !DllCall("WriteFile","uint",pipe,"str",Script,"uint",StrLen(Script)+1,"uint*",0,"uint",0)
      MsgBox WriteFile failed: %ErrorLevel%/%A_LastError%
   DllCall("CloseHandle","uint",pipe)
   Return PID
}

;by Lexikos: http://www.autohotkey.com/forum/viewtopic.php?t=25867
CreateNamedPipe(Name, OpenMode=3, PipeMode=0, MaxInstances=255)
{
    return DllCall("CreateNamedPipe","str","\\.\pipe\" Name,"uint",OpenMode
        ,"uint",PipeMode,"uint",MaxInstances,"uint",0,"uint",0,"uint",0,"uint",0)
}

;By SKAN: http://www.autohotkey.com/forum/viewtopic.php?t=44873
;modifications (lines with ;;): added 1 and 2, and replaced TrayTip... with Return 3
USBD_SafelyRemove( Drv ) {
 If A_OSVersion not in  WIN_VISTA,WIN_XP,WIN_2000
   Return 1 ;;
 If ! ( Serial := USBD_GetDeviceSerial( Drv ) )
   Return 2 ;;
 DeviceID := USBD_GetDeviceID( Serial )
 DeviceEject( DeviceID )
 IfExist, %Drv%\, Return 3 ;; TrayTip, %DeviceID%, Drive %Drv% was not Ejected!, 10, 3
 Else, TrayTip, %DeviceID%, Drive %Drv% was safely Removed, 10, 1
}

USBD_GetDeviceSerial( Drv="" ) {
 DriveGet, DriveType, Type, %Drv%
 IfNotEqual,DriveType,Removable, Return
 RegRead, Hex, HKLM, SYSTEM\MountedDevices, \DosDevices\%Drv%
 VarSetCapacity(U,(Sz:=StrLen(Hex)//2)),  VarSetCapacity(A,Sz+1)
 Loop % Sz
  NumPut( "0x" . SubStr(hex,2*A_Index-1,2), U, A_Index-1, "Char" )
 DllCall( "WideCharToMultiByte", Int,0,Int,0, UInt,&U,UInt,Sz, Str,A,UInt,Sz, Int,0,Int,0)
 StringSplit, Part, A, #
 ParentIdPrefixCheck := SubStr( Part3,1,InStr(Part3,"&",0,0)-1 )
 IfEqual,A_OSVersion,WIN_VISTA, Return,ParentIdPrefixCheck
 Loop, HKLM, SYSTEM\CurrentControlSet\Enum\USBSTOR,1,0
  { Device := A_LoopRegName
    Loop, HKLM, SYSTEM\CurrentControlSet\Enum\USBSTOR\%Device%,1,0
     { Serial := A_LoopRegName
       RegRead, PIPrefix, HKLM, SYSTEM\CurrentControlSet\Enum\USBSTOR\%Device%\%Serial%
              , ParentIdPrefix
       If ( PIPrefix = ParentIdPrefixCheck )
         Return, SubStr( Serial,1,InStr(Serial,"&",0,0)-1 )
     }
}}

USBD_GetDeviceID( Serial ) {
 Loop, HKLM, SYSTEM\CurrentControlSet\Enum\USB\,1,0
  { Device := A_LoopRegName
    Loop, HKLM, SYSTEM\CurrentControlSet\Enum\USB\%Device%,1,0
    If ( A_LoopRegName=Serial )
      Return DllCall( "CharUpperA", Str, "USB\" Device "\" Serial, Str )
}}

DeviceEject( DeviceID ) {
 hMod := DllCall( "LoadLibrary", Str,"SetupAPI.dll" ), VarSetCapacity(VE,255,0)
 If ! DllCall( "SetupAPI\CM_Locate_DevNodeA", UIntP,DI, Str,DeviceID, Int,0 )
 If ! DllCall( "SetupAPI\CM_Get_DevNode_Status", UIntP,STS, UIntP,PR, UInt,DI, Int,0)
 DllCall( "SetupAPI\CM_Request_Device_EjectA", UInt,DI, UIntP,VT, Str,VE, UInt,255, Int,0)
 DllCall( "FreeLibrary", UInt,hMod )
}