A Python-ctypes script to lock and unlock a cd tray by dispatching IOCTL in Windows

lock_cd_tray.py

# The MIT License (MIT)
#
# Copyright © 2020 Brian Maissy <[email protected]>
#
# Based on https://gist.github.com/santa4nt/11068180
# Copyright © 2014-2016 Santoso Wijaya <[email protected]>
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sub-license, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.


import argparse
import ctypes
import ctypes.wintypes as wintypes
from ctypes import windll, GetLastError, FormatError

LPDWORD = ctypes.POINTER(wintypes.DWORD)
LPOVERLAPPED = wintypes.LPVOID
LPSECURITY_ATTRIBUTES = wintypes.LPVOID

GENERIC_READ = 0x80000000
GENERIC_WRITE = 0x40000000

CREATE_NEW = 1
CREATE_ALWAYS = 2
OPEN_EXISTING = 3

FILE_ATTRIBUTE_NORMAL = 0x00000080

INVALID_HANDLE_VALUE = -1

NULL = 0


def _CreateFile(filename, access, mode, creation, flags):
    """CreateFile function
    http://msdn.microsoft.com/en-us/library/windows/desktop/aa363858(v=vs.85).aspx
    """
    CreateFile_Fn = windll.kernel32.CreateFileW
    CreateFile_Fn.argtypes = [
        wintypes.LPWSTR,        # _In_          LPCTSTR lpFileName
        wintypes.DWORD,         # _In_          DWORD dwDesiredAccess
        wintypes.DWORD,         # _In_          DWORD dwShareMode
        LPSECURITY_ATTRIBUTES,  # _In_opt_      LPSECURITY_ATTRIBUTES lpSecurityAttributes
        wintypes.DWORD,         # _In_          DWORD dwCreationDisposition
        wintypes.DWORD,         # _In_          DWORD dwFlagsAndAttributes
        wintypes.HANDLE         # _In_opt_      HANDLE hTemplateFile
    ]
    CreateFile_Fn.restype = wintypes.HANDLE

    return wintypes.HANDLE(CreateFile_Fn(filename, access, mode, NULL, creation, flags, NULL))


def _DeviceIoControl(devhandle, ioctl, inbuf, inbufsiz, outbuf, outbufsiz):
    """DeviceIoControl function
    http://msdn.microsoft.com/en-us/library/aa363216(v=vs.85).aspx
    """
    DeviceIoControl_Fn = windll.kernel32.DeviceIoControl
    DeviceIoControl_Fn.argtypes = [
        wintypes.HANDLE,    # _In_          HANDLE hDevice
        wintypes.DWORD,     # _In_          DWORD dwIoControlCode
        wintypes.LPVOID,    # _In_opt_      LPVOID lpInBuffer
        wintypes.DWORD,     # _In_          DWORD nInBufferSize
        wintypes.LPVOID,    # _Out_opt_     LPVOID lpOutBuffer
        wintypes.DWORD,     # _In_          DWORD nOutBufferSize
        LPDWORD,            # _Out_opt_     LPDWORD lpBytesReturned
        LPOVERLAPPED        # _Inout_opt_   LPOVERLAPPED lpOverlapped
    ]
    DeviceIoControl_Fn.restype = wintypes.BOOL

    # allocate a DWORD, and take its reference
    dwBytesReturned = wintypes.DWORD(0)
    lpBytesReturned = ctypes.byref(dwBytesReturned)

    status = DeviceIoControl_Fn(devhandle, ioctl, inbuf, inbufsiz, outbuf, outbufsiz, lpBytesReturned, None)

    return status, dwBytesReturned


class HandleError(Exception):
    pass


def format_last_error():
    error_code = GetLastError()
    return f'{FormatError(error_code)} (error {error_code})'


class DeviceIoControl(object):
    def __init__(self, path):
        self.path = path
        self._file_handle = None

    def _validate_handle(self):
        if self._file_handle is None:
            raise HandleError('No file handle')
        if self._file_handle.value == wintypes.HANDLE(INVALID_HANDLE_VALUE).value:
            raise HandleError(f'Failed to open {self.path}. {format_last_error()}')

    def ioctl(self, ctl, inbuf, inbufsiz, outbuf, outbufsiz):
        self._validate_handle()
        return _DeviceIoControl(self._file_handle, ctl, inbuf, inbufsiz, outbuf, outbufsiz)

    def __enter__(self):
        self._file_handle = _CreateFile(self.path, GENERIC_READ | GENERIC_WRITE, 0, OPEN_EXISTING,
                                        FILE_ATTRIBUTE_NORMAL)
        self._validate_handle()
        return self

    def __exit__(self, typ, val, tb):
        try:
            self._validate_handle()
        except HandleError:
            pass
        else:
            windll.kernel32.CloseHandle(self._file_handle)


# see https://docs.microsoft.com/en-us/windows/win32/api/winioctl/ni-winioctl-ioctl_storage_media_removal
# number from http://www.ioctls.net
IOCTL_STORAGE_MEDIA_REMOVAL = 0x2d4804


class PREVENT_MEDIA_REMOVAL(ctypes.Structure):
    """See: https://docs.microsoft.com/en-us/windows/win32/api/winioctl/ns-winioctl-prevent_media_removal"""
    _fields_ = [
        ('PreventMediaRemoval', wintypes.BOOLEAN),
    ]


def main():
    argument_parser = argparse.ArgumentParser()
    argument_parser.add_argument('operation', choices=['lock', 'unlock'])
    argument_parser.add_argument('--drive-letter', '-d', default='D')
    arguments = argument_parser.parse_args()

    value = wintypes.BOOLEAN(1 if arguments.operation == 'lock' else 0)
    prevent_removal = PREVENT_MEDIA_REMOVAL(value)

    with DeviceIoControl(r'\\.\{}:'.format(arguments.drive_letter)) as dctl:
        status, _ = dctl.ioctl(
            IOCTL_STORAGE_MEDIA_REMOVAL,
            ctypes.pointer(prevent_removal), ctypes.sizeof(prevent_removal),
            None, 0  # no output buffer
        )

    if status:
        print(f'{arguments.operation}ed successfully')
        return 0
    else:
        print(f'IOCTL returned failure. {format_last_error()}')
        return GetLastError()


if __name__ == '__main__':
    exit(main())