Improve PawnIO discovery and formatting changes

Author: Steve-Tech

cros_ec_python/devices/pawnio.py

@@ -5,6 +5,8 @@
 import sys
 from ctypes import wintypes
 from ctypes import util as ctypes_util
+import winreg
+from typing import Iterable
 
 from ..baseclass import CrosEcClass
 from ..constants.COMMON import *
@@ -28,11 +30,9 @@ def __init__(self, dll: str | None = None, bin: str | None = None):
         else:
             self.bin = bin or "LpcCrOSEC.bin"
 
-        if dll or (dll := ctypes_util.find_library("PawnIOLib.dll")):
-            self.pawniolib = ctypes.OleDLL(dll)
-        else:
-            # Let this raise an error if we can't find it
-            self.pawniolib = ctypes.OleDLL("C:\\Program Files\\PawnIO\\PawnIOLib.dll")
+        if not dll:
+            dll = self._get_location()
+        self.pawniolib = ctypes.OleDLL(dll)
 
         self.pawniolib.pawnio_version.argtypes = [ctypes.POINTER(wintypes.ULONG)]
         self.pawniolib.pawnio_open.argtypes = [ctypes.POINTER(wintypes.HANDLE)]
@@ -81,8 +81,8 @@ def _pawnio_load(self, filepath: str) -> None:
     def _pawnio_execute(
         self,
         function: str,
-        in_data: bytes,
-        out_size: bytes,
+        in_data: Iterable,
+        out_size: int,
         in_size: int | None = None,
     ) -> tuple[int | ctypes.Array]:
         function_bytes = function.encode("utf-8")
@@ -106,14 +106,29 @@ def _pawnio_execute(
     def _pawnio_close(self):
         self.pawniolib.pawnio_close(self.handle)
 
+    @staticmethod
+    def _get_location() -> str | None:
+        """
+        Get the location of the PawnIO driver.
+        """
+        try:
+            with winreg.OpenKey(
+                winreg.HKEY_LOCAL_MACHINE,
+                r"SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall\PawnIO"
+            ) as key:
+                return winreg.QueryValueEx(key, "InstallLocation")[0] + r"\\PawnIOLib.dll"
+        except FileNotFoundError:
+            fallback = os.environ.get("ProgramFiles", r"C:\Program Files") + r"\PawnIO\PawnIOLib.dll"
+            if os.path.exists(fallback):
+                return fallback
+            return None
+
     @staticmethod
     def detect() -> bool:
         """
         Detect if the PawnIO driver is installed.
         """
-        return bool(ctypes_util.find_library("PawnIOLib.dll")) or os.path.exists(
-            "C:\\Program Files\\PawnIO\\PawnIOLib.dll"
-        )
+        return CrosEcPawnIO._get_location() is not None
 
     def ec_init(self) -> None:
         self._pawnio_open()

cros_ec_python/devices/win_fw_ec.py

@@ -29,29 +29,31 @@
 INVALID_HANDLE_VALUE: Final = -1
 
 
-def ctl_code(device_type, function, method, access) -> int:
-    return ((device_type) << 16) + ((access) << 14) + ((function) << 2) + method
+def CTL_CODE(device_type, function, method, access) -> int:
+    return (device_type << 16) | (access << 14) | (function << 2) | method
 
 
 CROSEC_CMD_MAX_REQUEST = 0x100
 
 FILE_DEVICE_CROS_EMBEDDED_CONTROLLER = 0x80EC
 
-IOCTL_CROSEC_XCMD = ctl_code(
+IOCTL_CROSEC_XCMD = CTL_CODE(
     FILE_DEVICE_CROS_EMBEDDED_CONTROLLER,
     0x801,
     METHOD_BUFFERED,
     FILE_READ_DATA | FILE_WRITE_DATA,
 )
-IOCTL_CROSEC_RDMEM = ctl_code(
+IOCTL_CROSEC_RDMEM = CTL_CODE(
     FILE_DEVICE_CROS_EMBEDDED_CONTROLLER,
     0x802,
     METHOD_BUFFERED,
     FILE_READ_ACCESS,
 )
 
 
-def CreateFileW(filename, access, mode, creation, flags) -> wintypes.HANDLE:
+def CreateFileW(
+    filename: str, access: int, mode: int, creation: int, flags: int
+) -> wintypes.HANDLE:
     knl32_CreateFileW = windll.kernel32.CreateFileW
     knl32_CreateFileW.argtypes = [
             wintypes.LPCWSTR, # lpFileName
@@ -69,7 +71,14 @@ def CreateFileW(filename, access, mode, creation, flags) -> wintypes.HANDLE:
     )
 
 
-def DeviceIoControl(devhandle, ioctl, inbuf, inbufsiz, outbuf, outbufsiz) -> bool:
+def DeviceIoControl(
+    handle: wintypes.HANDLE,
+    ioctl: int,
+    inbuf: ctypes.pointer,
+    insize: int,
+    outbuf: ctypes.pointer,
+    outsize,
+) -> bool:
     knl32_DeviceIoControl = windll.kernel32.DeviceIoControl
     knl32_DeviceIoControl.argtypes = [
         wintypes.HANDLE,  # hDevice
@@ -84,7 +93,7 @@ def DeviceIoControl(devhandle, ioctl, inbuf, inbufsiz, outbuf, outbufsiz) -> boo
     knl32_DeviceIoControl.restype = wintypes.BOOL
 
     status = knl32_DeviceIoControl(
-        devhandle, ioctl, inbuf, inbufsiz, outbuf, outbufsiz, None, None
+        handle, ioctl, inbuf, insize, outbuf, outsize, None, None
     )
 
     return bool(status)