AutoHotkey Community

; FlashLock.ahk
; ~~~~~~~~~~~~~
; Lock/Unlock the workstation with a USB device of choice.
; Right click on tray icon and select "Set Device" to add it.
; Add dword HKCU\Software\FlashLock\REG_DIM_SCREEN_VAL to change background transparency.
; ----------------------------------
; TODO: 1. Block CTRL+ALT+DEL
;       2. Enumerate devices on start and check if the USB key is inserted.
; ----------------------------------
; cyrusza - http://ciroprincipe.info

#SingleInstance force
#Persistent
OnExit, QUIT

; ############################################
; ############### MAIN SECTION ###############
; ############################################

Menu, TRAY, NoStandard
Menu, TRAY, Add, Set Device, DEVSET
Menu, TRAY, Add
Menu, TRAY, Add, Quit, QUIT

IconDataHex =
( Join
0000010001001010000000002000680400001600000028000000100000002000000001002000000000004004000000000000000000000000000000000000ffff
ff01369dd9ff3199d8ff2c94d7ff2890d6ff238cd5ff1e88d4ff1a84d3ff1580d2ff117cd1ff0e79d1ff0a76d0ff0773cfff0470cfff016eceffffffff01ffff
ff013da3daffbcebfaffbcebfcffbfeefeffc6f4ffffcef8ffffd3faffffd0f8ffffc7f2ffffbae9fcffb3e4f9ffb0e2f8ffb0e2f8ff0571cfffffffff01ffff
ff0143a8dbffbfecfbff59cff5ff41b0ecff4ebaefff5ac2efff60c6efff5cc4efff4cb6efff37a5e6ff2a9ae1ff38b8eeffb1e3f8ff0975d0ffffffff01ffff
ff0149addcffc1eefbff5fd3f7ff6cdbfcff7fe5ffff8fedffff97f2ffff93edffff7cdfffff5bccf8ff46beefff3cbaeeffb3e3f9ff0e79d1ffffffff01ffff
ff014eb2ddffc3effbff65d6f8ff4cb6ecff5abdefff95ebffff3097ddff4d82abff84e1ffff41a9e9ff329fe1ff42beefffb4e5f9ff137ed2ffffffff01ffff
ff0153b7deffc6f0fcff6ad9f8ff7ce2fdff90e8ffff99e9ffff329fdfff548bb2ff8ae2ffff6ad0f9ff50c5f1ff46c1f0ffb6e7f9ff1883d3ffffffff01ffff
ff0158bbdfffc7f1fcff6fdcf9ff56bbedff61bdefff9be7ffff35a6e2ff4ba4e1ff90e2ffff49ade9ff38a4e3ff49c4f0ffb8e8f9ff1e88d4ffffffff01ffff
ff015cbfe0ffc8f3fcff75dff9ff89e6fdff95e7ffff9ae5ffffaaeeffffa8edffff99e3ffff74d5f9ff59ccf3ff4fc8f1ffbbe9faff248dd5ffffffff01ffff
ff0160c2e1ffc9f3fcffcbf3fdffd4f6feffd7f6ffffd8f4ffffe0f8ffffdff8ffffdaf5ffffcdf1fcffc2edfaffbdebfaffbdebfaff2b93d6ffffffff01ffff
ff0161c3e1ff88a0a8ff919191ff8e8e8eff5ab9dcff55b8dfff51b5deff4db1ddff49addcff46a8d7ff787878ff767676ff657e8dff3199d8ffffffff01ffff
ff01ffffff01999999c3c6c6c6ff949494ff8f8f8f09ffffff01ffffff01ffffff01ffffff017e7e7e097d7d7dffabababff767676c3ffffff01ffffff01ffff
ff01ffffff019d9d9dafc4c4c4ffa1a1a1ff93939329ffffff01ffffff01ffffff01ffffff0182828229898989ffa9a9a9ff797979afffffff01ffffff01ffff
ff01ffffff01a1a1a175bababaffbfbfbfff989898dd9494941f919191058e8e8e058a8a8a1f878787e5a8a8a8ff9e9e9eff7d7d7d75ffffff01ffffff01ffff
ff01ffffff01a4a4a40da1a1a1dbc4c4c4ffbebebeffa1a1a1ff969696ff939393ff979797ffaeaeaeffaeaeaeff848484db81818109ffffff01ffffff01ffff
ff01ffffff01ffffff01a4a4a431a2a2a2dfbcbcbcffcacacaffccccccffcacacaffc2c2c2ffadadadff8c8c8cdf89898931ffffff01ffffff01ffffff01ffff
ff01ffffff01ffffff01ffffff01a5a5a50da3a3a381a0a0a0bb9d9d9dcd9a9a9acd979797bb949494819090900dffffff01ffffff01ffffff01ffffff010000
ffff0000ffff0000ffff0000ffff0000ffff0000ffff0000ffff0000ffff0000ffff0000ffff0000ffff0000ffff0000ffff0000ffff0000ffff0000ffff
)
SetTrayIcon(IconDataHex)

Gui +LastFound
DetectHiddenWindows, On
hMainWin := WinExist(A_ScriptFullPath " ahk_class AutoHotkey ahk_pid " DllCall("GetCurrentProcessId"))
DetectHiddenWindows, Off
OnMessage(WM_DEVICECHANGE:=0x219, "USBHandler")
hDevNotify := RegForDevIntNotify(GUID_DEVINTERFACE_USB_DEVICE:="{a5dcbf10-6530-11d2-901f-00c04fb951ed}", hMainWin)

isSetMode := 0, isLocked := 0
RegRead, dimScreenVal, HKCU, Software\FlashLock, REG_DIM_SCREEN_VAL
RegRead, setDeviceId, HKCU, Software\FlashLock, REG_SET_DEVICE_ID
If (ErrorLevel) {
  ; Progress SKAN style - http://www.autohotkey.com/forum/viewtopic.php?t=45718
  Progress, CWE6E3E4 CT000020 CBF73D00 w370 h46 B1 FS8 WM700 WS400 FM8 ZH0, Locking device required`, use the "Set Device" tray menu., FlashLock, FlashLock
  WinSet, Transparent, 210, FlashLock
  Sleep, 10000
  Progress, Off
}
Else {
  ToggleLock(1)
  ToggleBackground(1)
}

Return

; ############################################
; ############## LABELS SECTION ##############
; ############################################

DEVSET:
  isSetMode := 1
  Progress, CWE6E3E4 CT000020 CBF73D00 w370 h46 B1 FS8 WM700 WS400 FM8 ZH8 ZY3, , (Re)Insert the USB locking device, FlashLock
  WinSet, Transparent, 210, FlashLock
  Loop 100 {
    If (isSetMode = 0) { ; Set by USBHandler in case of USB insertion during the timeframe
      RegWrite, REG_SZ, HKCU, Software\FlashLock, REG_SET_DEVICE_ID, %setDeviceId%
      Progress, CWE6E3E4 CT000020 CBF73D00 w370 h46 B1 FS8 WM700 WS400 FM8 ZH0, Restarting program..., USB device set, FlashLock
      WinSet, Transparent, 210, FlashLock
      Sleep, 3000
      Reload
    }
    Progress, %A_Index%, % "Time remaining: " (100-A_Index)//5 " seconds"
    Sleep, 200
  }
  Progress, Off
  isSetMode := 0
  Return
; --- END DEVSET ---

QUIT:
  ; UnregisterDeviceNotification
  ; http://msdn.microsoft.com/en-us/library/windows/desktop/aa363475(v=vs.85).aspx
  DllCall("User32.dll\UnregisterDeviceNotification", UInt, hDevNotify)
  ExitApp
; --- END QUIT ---

DUMMY:
  Return
; --- END DUMMY ---

; ############################################
; ######### HELPER FUNCTIONS SECTION #########
; ############################################

; Registers for device notifications.
RegForDevIntNotify(DEVINT_GUID, hMainWin) {
  VarSetCapacity(sDBDI, 32, 0)    ; DEV_BROADCAST_DEVICEINTERFACE: 4 + 4 + 4 + 16 [4+2+2+1*8] + 1 + 3 [PAD]
  NumPut(32, sDBDI, 0, "UInt")
  NumPut((DBT_DEVTYP_DEVICEINTERFACE:=0x00000005), sDBDI, 4, "UInt")

  ; Put the GUID inside the DEV_BROADCAST_DEVICEINTERFACE structure.
  DEVINT_GUID := RegExReplace(DEVINT_GUID, "[{}-]", "", 0, 5)
  NumPut("0x" . SubStr(DEVINT_GUID, 1, 8), sDBDI, 12, "UInt")
  NumPut("0x" . SubStr(DEVINT_GUID, 9, 4), sDBDI, 16, "UShort")
  NumPut("0x" . SubStr(DEVINT_GUID, 13, 4), sDBDI, 18, "UShort")
  Loop, 8
    NumPut("0x" . SubStr(DEVINT_GUID, 15+A_Index*2, 2), sDBDI, 19+A_Index, "UChar")
 
  ; RegisterDeviceNotification
  ; http://msdn.microsoft.com/en-us/library/windows/desktop/aa363431(v=vs.85).aspx
  Return DllCall( "User32.dll\RegisterDeviceNotificationA"
                  , UInt, hMainWin
                  , UInt, &sDBDI
                  , UInt, 0x00000000)
}

; Handles and dispatches notifications.
USBHandler(wParam, lParam, msg, hWnd) {
  Global isSetMode, setDeviceId, hMainWin
  If (hWnd = hMainWin && NumGet(lParam+0, 4, "UInt") = 0x05) {  ; DBT_DEVTYP_DEVICEINTERFACE = 0x05
    Addr := lParam + 28
    While (aChar := *Addr++) != 0
      devId .= Chr(aChar)
    RegExMatch(devId, "(?<=\#).*(?=\#)", devId)
   
    If (wParam = 0x8000) {                              ; DBT_DEVICEARRIVAL = 0x8000
      If (isSetMode) {
        setDeviceId := devId
        isSetMode := 0
      }
      Else If (devId = setDeviceId) {
        ToggleLock(0)
        ToggleBackground(0)
      }
    }
    Else If (wParam = 0x8004 && devId = setDeviceId) {  ; DBT_DEVICEREMOVECOMPLETE = 0x8004
      ToggleLock(1)
      ToggleBackground(1)
    }
  }
}

; Locks/Unlocks keystrokes and mouse clicks.
ToggleLock(cmd) {
  Global isLocked
  SetFormat, IntegerFast, Hex
  If (cmd) {
    isLocked := 1
    Loop 0x1FF  ; Blocks keystrokes
      Hotkey, sc%A_Index%, DUMMY, On
    Loop 0x7    ; Blocks mouse clicks
      Hotkey, *vk%A_Index%, DUMMY, On
  }
  Else {
    isLocked := 0
    Loop 0x1FF  ; Unblocks keystrokes
      Hotkey, sc%A_Index%, DUMMY, Off
    Loop 0x7    ; Unblocks mouse clicks
      Hotkey, *vk%A_Index%, DUMMY, Off
  }
  SetFormat, Integer, D
}

; Turns On/Off the background and the front message.
ToggleBackground(cmd) {
  Global dimScreenVal
  If (cmd) {
    ; Creates the dark background
    Gui, 11:+LastFound +AlwaysOnTop +ToolWindow -Caption
    Gui, 11:Color, 000000, 000000
    WinSet, Transparent, % (dimScreenVal) ? dimScreenVal : "180"
    Gui, 11:Show, W%A_ScreenWidth% H%A_ScreenHeight% X0 Y0
    ; Turn on the front message
    Progress, CWE6E3E4 CT000020 CBF73D00 w370 h46 B1 FS8 WM700 WS400 FM8 ZH0, Waiting for the USB unlocking device., Workstation locked, FlashLock
    WinSet, Transparent, 210, FlashLock
  }
  Else {
    ; Turns off the front message
    Progress, Off
    ; Destroys the dark background
    Gui, 11:Destroy
  }
}

; Thanks SKAN               - http://www.autohotkey.com/forum/viewtopic.php?t=33955
SetTrayIcon(ByRef IconDataHex) {
  VarSetCapacity(IconData, (nSize:=StrLen(IconDataHex)//2))
  Loop %nSize% ; MCode by Laszlo Hars: http://www.autohotkey.com/forum/viewtopic.php?t=21172
    NumPut("0x" . SubStr(IconDataHex, 2*A_Index - 1, 2), IconData, A_Index-1, "UChar")
  IconDataHex := ""

  hICon := DllCall( "User32.dll\CreateIconFromResourceEx"
                    , UInt, &IconData+22
                    , UInt, NumGet(IconData, 14)
                    , Int, 1
                    , UInt, 0x00030000
                    , Int, 16
                    , Int, 16
                    , UInt, 0 )

  ; Thanks Chris            - http://www.autohotkey.com/forum/viewtopic.php?p=69461#69461
  Gui +LastFound
  SendMessage, (WM_SETICON:=0x80), 0, hIcon ; Set the Titlebar Icon
  SendMessage, (WM_SETICON:=0x80), 1, hIcon ; Set the Alt-Tab icon

  ; Creating NOTIFYICONDATA - http://www.msdn.microsoft.com/en-us/library/aa930660.aspx
  ; Thanks Lexikos          - http://www.autohotkey.com/forum/viewtopic.php?p=162175#162175
  VarSetCapacity(NID, 444, 0)
  NumPut(444, NID, 0)
  DetectHiddenWindows, On
  NumPut(WinExist(A_ScriptFullPath " ahk_class AutoHotkey ahk_pid " DllCall("GetCurrentProcessId")), NID, 4)
  DetectHiddenWindows, Off
  NumPut(1028, NID, 8)
  NumPut(2, NID, 12)
  NumPut(hIcon, NID, 20)

  Return DllCall( "Shell32.dll\Shell_NotifyIcon"
                  , UInt, 0x1
                  , UInt, &NID )
}