Files
sphaira/tools/usb_common.py
AndrewSpangler 0a8bc01870 Clean up python (#212)
* Python cleanup
2025-08-31 03:48:37 +01:00

145 lines
4.4 KiB
Python

import struct
from time import sleep
from usb.core
import usb.util
SPLASH = """
:@@@@@@@@@@@@@@@@@@@@@@@@@@:
#@ @#
#@ @#
#@ @#
#@ @@@@@@ @@@@@@ @#
#@ @@@@@@ @@@@@@ @#
#@ @@@@@@ @@@@@@ @#
#@ @#
#@ @#
#@ @#
@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@ @% @@@@@@@@@@@@
#@ *@@@@* @% @@@@@@@@@@@@@
#@ @@@@@@@@ @% @@@@@@@@@@@@@
#@ @@@@@@@@ @% @@@@@@@@@@@@@
#@ *@@@@* @% @@@@@@@@@@@@@
#@ @% @@@@@@@@@@@@@
#@ @% @@@@@@@@@@@@@
#@ @% @@@@@@@@@@@@@
#@ @% @@@@@@@@@@@@@
#@ @% @@@@@@@@@@@@@
#@ @% @@@@@=--=@@@@
#@ @% @@@- -@@
#@ @% @@@ @@
#@ @% @@@- -@@
=@@ @% @@@@@=--=@@@@
=@@ @% @@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
"""
class USB_ENUM:
# (SPH0) for the script and Switch
MAGIC = 0x53504830
# Commands
CMD_QUIT = 0
CMD_OPEN = 1
CMD_EXPORT = 1
# Result Codes
RESULT_OK = 0
RESULT_ERROR = 1
# Flags
FLAG_NONE = 0
FLAG_STREAM = 1 << 0
# Switch Vendor / Product ID
VENDOR_ID = 0x057E
PRODUCT_ID = 0x3000
ENABLE_ZLT = 0
def find_switch() -> object | None:
return usb.core.find(
idVendor=USB_ENUM.VENDOR_ID,
idProduct=USB_ENUM.PRODUCT_ID
)
class Usb:
def __init__(self):
self._out_ep = None
self._in_ep = None
self._packet_size = 0
self._packet_index = 0
def wait_for_connect(self) -> None:
print(SPLASH)
print("Waiting for Switch...", end="")
dev = None
while dev is None:
if (dev := find_switch()):
break
print(".", end="")
sleep(0.5)
print("Found the Switch!\n")
cfg = None
print("Getting configuration...")
try:
cfg = dev.get_active_configuration()
print("Found active config")
except usb.core.USBError:
print("No currently active config")
cfg = None
if cfg is None:
dev.reset()
dev.set_configuration()
cfg = dev.get_active_configuration()
is_out_ep = lambda ep: usb.util.endpoint_direction(ep.bEndpointAddress) == usb.util.ENDPOINT_OUT
is_in_ep = lambda ep: usb.util.endpoint_direction(ep.bEndpointAddress) == usb.util.ENDPOINT_IN
self._out_ep = usb.util.find_descriptor(cfg[(0,0)], custom_match=is_out_ep)
self._in_ep = usb.util.find_descriptor(cfg[(0,0)], custom_match=is_in_ep)
if not self._out_ep:
raise ValueError("Failed to get USB OUT address")
if not self._in_ep:
raise ValueError("Failed to get USB IN address")
print(f"iManufacturer: {dev.manufacturer} \
iProduct: {dev.product} \
iSerialNumber: {dev.serial_number}")
print(f"bcdUSB: {hex(dev.bcdUSB)} \
bMaxPacketSize0: {dev.bMaxPacketSize0}")
self._packet_size = 1 << dev.bMaxPacketSize0
def read(self, size: int, timeout: int = 0) -> bytes:
if (ENABLE_ZLT and size and not (size % self._packet_size)):
size += 1
return self._in_ep.read(size, timeout)
def write(self, buf: bytes, timeout: int = 0) -> int:
packet = self._packet_index
self._packet_index += 1
# Todo, implement packet index long
return self._out_ep.write(data=buf, timeout=timeout)
def get_send_header(self) -> tuple[int, int, int]:
header = self.read(16)
magic, arg2, arg3, arg4 = struct.unpack('<IIII', header)
if magic != MAGIC:
raise Exception("Unexpected magic {}".format(magic))
return arg2, arg3, arg4
def get_send_data_header(self) -> tuple[int, int, int]:
header = self.read(16)
return struct.unpack('<QII', header)
def send_result(self, result: int, arg3: int = 0, arg4: int = 0) -> None:
send_data = struct.pack('<IIII', MAGIC, result, arg3, arg4)
self.write(send_data)