refactor: Add Python type annotations wherever appropriate (#269)

* Add Python type annotations wherever appropriate

* Might as well annotate this too

proximity_keys.py

@@ -4,50 +4,53 @@
 # See https://github.com/google/bumble/blob/main/docs/mkdocs/src/platforms/windows.md for usage.
 # You need to associate WinUSB with your Bluetooth interface. Once done, you can roll back to the original driver from Device Manager.
 
-import sys
 import asyncio
-import argparse
+import colorama
 import logging
 import platform
-from typing import Any, Optional
+from argparse import ArgumentParser, Namespace
+from asyncio import Queue, TimeoutError
+from colorama import Fore, Style
+from logging import Formatter, LogRecord, Logger, StreamHandler
+from socket import socket as Socket
+from typing import Any, Dict, List, Optional, Tuple
 
-from colorama import Fore, Style, init as colorama_init
-colorama_init(autoreset=True)
+colorama.init(autoreset=True)
 
-handler = logging.StreamHandler()
-class ColorFormatter(logging.Formatter):
-    COLORS = {
+handler: StreamHandler = StreamHandler()
+class ColorFormatter(Formatter):
+    COLORS: Dict[int, str] = {
         logging.DEBUG: Fore.BLUE,
         logging.INFO: Fore.GREEN,
         logging.WARNING: Fore.YELLOW,
         logging.ERROR: Fore.RED,
         logging.CRITICAL: Fore.MAGENTA,
     }
-    def format(self, record):
-        color = self.COLORS.get(record.levelno, "")
-        prefix = f"{color}[{record.levelname}:{record.name}]{Style.RESET_ALL}"
+    def format(self, record: LogRecord) -> str:
+        color: str = self.COLORS.get(record.levelno, "")
+        prefix: str = f"{color}[{record.levelname}:{record.name}]{Style.RESET_ALL}"
         return f"{prefix} {record.getMessage()}"
 handler.setFormatter(ColorFormatter())
 logging.basicConfig(level=logging.INFO, handlers=[handler])
-logger = logging.getLogger("proximitykeys")
+logger: Logger = logging.getLogger("proximitykeys")
 
-PROXIMITY_KEY_TYPES = {0x01: "IRK", 0x04: "ENC_KEY"}
+PROXIMITY_KEY_TYPES: Dict[int, str]  = {0x01: "IRK", 0x04: "ENC_KEY"}
 
-def parse_proximity_keys_response(data: bytes):
+def parse_proximity_keys_response(data: bytes) -> Optional[List[Tuple[str, bytes]]]:
     if len(data) < 7 or data[4] != 0x31:
         return None
-    key_count = data[6]
-    keys = []
-    offset = 7
+    key_count: int = data[6]
+    keys: List[Tuple[str, bytes]] = []
+    offset: int = 7
     for _ in range(key_count):
         if offset + 3 >= len(data):
             break
-        key_type = data[offset]
-        key_length = data[offset + 2]
+        key_type: int = data[offset]
+        key_length: int = data[offset + 2]
         offset += 4
         if offset + key_length > len(data):
             break
-        key_bytes = data[offset:offset + key_length]
+        key_bytes: bytes = data[offset:offset + key_length]
         keys.append((PROXIMITY_KEY_TYPES.get(key_type, f"TYPE_{key_type:02X}"), key_bytes))
         offset += key_length
     return keys
@@ -55,7 +58,7 @@ def parse_proximity_keys_response(data: bytes):
 def hexdump(data: bytes) -> str:
     return " ".join(f"{b:02X}" for b in data)
 
-async def run_bumble(bdaddr: str):
+async def run_bumble(bdaddr: str) -> int:
     try:
         from bumble.l2cap import ClassicChannelSpec
         from bumble.transport import open_transport
@@ -68,19 +71,23 @@ async def run_bumble(bdaddr: str):
         logger.error("Bumble not installed")
         return 1
 
-    PSM_PROXIMITY = 0x1001
-    HANDSHAKE = bytes.fromhex("00 00 04 00 01 00 02 00 00 00 00 00 00 00 00 00")
-    KEY_REQ = bytes.fromhex("04 00 04 00 30 00 05 00")
+    PSM_PROXIMITY: int = 0x1001
+    HANDSHAKE: bytes = bytes.fromhex("00 00 04 00 01 00 02 00 00 00 00 00 00 00 00 00")
+    KEY_REQ: bytes = bytes.fromhex("04 00 04 00 30 00 05 00")
 
     class KeyStore:
-        async def delete(self, name: str): pass
-        async def update(self, name: str, keys: Any): pass
-        async def get(self, _name: str) -> Optional[Any]: return None
-        async def get_all(self): return []
+        async def delete(self, name: str) -> None:
+            pass
+        async def update(self, name: str, keys: Any) -> None:
+            pass
+        async def get(self, _name: str) -> Optional[Any]:
+            return None
+        async def get_all(self) -> List[Tuple[str, Any]]:
+            return []
 
-        async def get_resolving_keys(self) -> list[tuple[bytes, Any]]:
-            all_keys = await self.get_all()
-            resolving_keys = []
+        async def get_resolving_keys(self) -> List[Tuple[bytes, Any]]:
+            all_keys: List[Tuple[str, Any]] = await self.get_all()
+            resolving_keys: List[Tuple[bytes, Any]] = []
             for name, keys in all_keys:
                 if getattr(keys, "irk", None) is not None:
                     resolving_keys.append((
@@ -89,8 +96,8 @@ async def run_bumble(bdaddr: str):
                     ))
             return resolving_keys
 
-    async def exchange_keys(channel, timeout=5.0):
-        recv_q: asyncio.Queue = asyncio.Queue()
+    async def exchange_keys(channel: Any, timeout: float = 5.0) -> Optional[List[Tuple[str, bytes]]]:
+        recv_q: Queue = Queue()
         channel.sink = lambda sdu: recv_q.put_nowait(sdu)
         logger.info("Sending handshake packet...")
         channel.send_pdu(HANDSHAKE)
@@ -99,19 +106,19 @@ async def run_bumble(bdaddr: str):
         channel.send_pdu(KEY_REQ)
         while True:
             try:
-                pkt = await asyncio.wait_for(recv_q.get(), timeout)
-            except asyncio.TimeoutError:
+                pkt: bytes = await asyncio.wait_for(recv_q.get(), timeout)
+            except TimeoutError:
                 logger.error("Timed out waiting for SDU response")
                 return None
             logger.debug("Received SDU (%d bytes): %s", len(pkt), hexdump(pkt))
-            keys = parse_proximity_keys_response(pkt)
+            keys: Optional[List[Tuple[str, bytes]]] = parse_proximity_keys_response(pkt)
             if keys:
                 return keys
 
-    async def get_device():
+    async def get_device() -> Tuple[Any, Device]:
         logger.info("Opening transport...")
-        transport = await open_transport("usb:0")
-        device = Device(host=Host(controller_source=transport.source, controller_sink=transport.sink))
+        transport: Any = await open_transport("usb:0")
+        device: Device = Device(host=Host(controller_source=transport.source, controller_sink=transport.sink))
         device.classic_enabled = True
         device.le_enabled = False
         device.keystore = KeyStore()
@@ -123,15 +130,15 @@ async def run_bumble(bdaddr: str):
         logger.info("Device powered on")
         return transport, device
 
-    async def create_channel_and_exchange(conn):
-        spec = ClassicChannelSpec(psm=PSM_PROXIMITY, mtu=2048)
+    async def create_channel_and_exchange(conn: Any) -> None:
+        spec: ClassicChannelSpec = ClassicChannelSpec(psm=PSM_PROXIMITY, mtu=2048)
         logger.info("Requesting L2CAP channel on PSM = 0x%04X", spec.psm)
         if not conn.is_encrypted:
             logger.info("Enabling link encryption...")
             await conn.encrypt()
             await asyncio.sleep(0.05)
-        channel = await conn.create_l2cap_channel(spec=spec)
-        keys = await exchange_keys(channel, timeout=8.0)
+        channel: Any = await conn.create_l2cap_channel(spec=spec)
+        keys: Optional[List[Tuple[str, bytes]]] = await exchange_keys(channel, timeout=8.0)
         if not keys:
             logger.warning("No proximity keys found")
             return
@@ -165,14 +172,14 @@ async def run_bumble(bdaddr: str):
         logger.info("Transport closed")
     return 0
 
-def run_linux(bdaddr: str):
+def run_linux(bdaddr: str) -> None:
     import socket
-    PSM = 0x1001
-    handshake = bytes.fromhex("00 00 04 00 01 00 02 00 00 00 00 00 00 00 00 00")
-    key_req = bytes.fromhex("04 00 04 00 30 00 05 00")
+    PSM: int = 0x1001
+    handshake: bytes = bytes.fromhex("00 00 04 00 01 00 02 00 00 00 00 00 00 00 00 00")
+    key_req: bytes = bytes.fromhex("04 00 04 00 30 00 05 00")
 
     logger.info("Connecting to %s (L2CAP)...", bdaddr)
-    sock = socket.socket(socket.AF_BLUETOOTH, socket.SOCK_SEQPACKET, socket.BTPROTO_L2CAP)
+    sock: Socket = Socket(socket.AF_BLUETOOTH, socket.SOCK_SEQPACKET, socket.BTPROTO_L2CAP)
     try:
         sock.connect((bdaddr, PSM))
         logger.info("Connected, sending handshake and key request...")
@@ -180,9 +187,9 @@ def run_linux(bdaddr: str):
         sock.send(key_req)
 
         while True:
-            pkt = sock.recv(1024)
+            pkt: bytes = sock.recv(1024)
             logger.debug("Received packet (%d bytes): %s", len(pkt), hexdump(pkt))
-            keys = parse_proximity_keys_response(pkt)
+            keys: Optional[List[Tuple[str, bytes]]] = parse_proximity_keys_response(pkt)
             if keys:
                 logger.info("Keys successfully retrieved")
                 print(f"{Fore.CYAN}{Style.BRIGHT}Proximity Keys:{Style.RESET_ALL}")
@@ -197,12 +204,12 @@ def run_linux(bdaddr: str):
         sock.close()
         logger.info("Connection closed")
 
-def main():
-    parser = argparse.ArgumentParser()
+def main() -> None:
+    parser: ArgumentParser = ArgumentParser()
     parser.add_argument("bdaddr")
     parser.add_argument("--debug", action="store_true")
     parser.add_argument("--bumble", action="store_true")
-    args = parser.parse_args()
+    args: Namespace = parser.parse_args()
     logging.getLogger().setLevel(logging.DEBUG if args.debug else logging.INFO)
 
     if args.bumble or platform.system() == "Windows":