From a42b104c354b6a1685415f42e5a73fa97fe66831 Mon Sep 17 00:00:00 2001
From: Adeeb Shihadeh <adeebshihadeh@gmail.com>
Date: Thu, 3 Nov 2022 22:18:12 -0700
Subject: [PATCH] SPI handle in python lib (#1131)

* SPI handle in python lib

* add spidev

* unify those

* clean up connect

* clean that up

Co-authored-by: Robbe Derks <robbe.derks@gmail.com>
---
 python/__init__.py   | 40 +++++++++++++++---------
 python/spi.py        | 74 ++++++++++++++++++++++++++++++++++++++++++++
 requirements.txt     |  1 +
 tests/health_test.py | 27 +++++++++-------
 4 files changed, 117 insertions(+), 25 deletions(-)
 create mode 100644 python/spi.py

diff --git a/python/__init__.py b/python/__init__.py
index ceb66627..41c11cc3 100644
--- a/python/__init__.py
+++ b/python/__init__.py
@@ -11,12 +11,14 @@ import warnings
 from functools import wraps
 from typing import Optional
 from itertools import accumulate
+
 from .dfu import PandaDFU, MCU_TYPE_F2, MCU_TYPE_F4, MCU_TYPE_H7  # pylint: disable=import-error
 from .flash_release import flash_release  # noqa pylint: disable=import-error
 from .update import ensure_st_up_to_date  # noqa pylint: disable=import-error
 from .serial import PandaSerial  # noqa pylint: disable=import-error
 from .isotp import isotp_send, isotp_recv  # pylint: disable=import-error
 from .config import DEFAULT_FW_FN, DEFAULT_H7_FW_FN, SECTOR_SIZES_FX, SECTOR_SIZES_H7  # noqa pylint: disable=import-error
+from .spi import SpiHandle # noqa pylint: disable=import-error
 
 __version__ = '0.0.10'
 
@@ -230,7 +232,7 @@ class Panda:
   FLAG_GM_HW_CAM = 1
   FLAG_GM_HW_CAM_LONG = 2
 
-  def __init__(self, serial: Optional[str] = None, claim: bool = True, disable_checks: bool = True):
+  def __init__(self, serial: Optional[str] = None, claim: bool = True, spi: bool = False, disable_checks: bool = True):
     self._serial = serial
     self._disable_checks = disable_checks
 
@@ -238,9 +240,9 @@ class Panda:
     self._bcd_device = None
 
     # connect and set mcu type
+    self._spi = spi
     self.connect(claim)
 
-
   def close(self):
     self._handle.close()
     self._handle = None
@@ -248,10 +250,30 @@ class Panda:
   def connect(self, claim=True, wait=False):
     if self._handle is not None:
       self.close()
-
-    context = usb1.USBContext()
     self._handle = None
 
+    if self._spi:
+      self._handle = SpiHandle()
+
+      # TODO implement
+      self._serial = "SPIDEV"
+      self.bootstub = False
+
+    else:
+      self.usb_connect(claim=claim, wait=wait)
+
+    assert self._handle is not None
+    self._mcu_type = self.get_mcu_type()
+    self.health_version, self.can_version, self.can_health_version = self.get_packets_versions()
+    print("connected")
+
+    # disable openpilot's heartbeat checks
+    if self._disable_checks:
+      self.set_heartbeat_disabled()
+      self.set_power_save(0)
+
+  def usb_connect(self, claim=True, wait=False):
+    context = usb1.USBContext()
     while 1:
       try:
         for device in context.getDeviceList(skip_on_error=True):
@@ -284,16 +306,6 @@ class Panda:
         break
       context = usb1.USBContext()  # New context needed so new devices show up
 
-    assert self._handle is not None
-    self._mcu_type = self.get_mcu_type()
-    self.health_version, self.can_version, self.can_health_version = self.get_packets_versions()
-    print("connected")
-
-    # disable openpilot's heartbeat checks
-    if self._disable_checks:
-      self.set_heartbeat_disabled()
-      self.set_power_save(0)
-
   def reset(self, enter_bootstub=False, enter_bootloader=False, reconnect=True):
     try:
       if enter_bootloader:
diff --git a/python/spi.py b/python/spi.py
new file mode 100644
index 00000000..a84c4c79
--- /dev/null
+++ b/python/spi.py
@@ -0,0 +1,74 @@
+import struct
+import spidev
+from functools import reduce
+
+# Constants
+SYNC = 0x5A
+HACK = 0x79
+DACK = 0x85
+NACK = 0x1F
+CHECKSUM_START = 0xAB
+
+MAX_RETRY_COUNT = 5
+
+# This mimics the handle given by libusb1 for easy interoperability
+class SpiHandle:
+  def __init__(self):
+    self.spi = spidev.SpiDev()
+    self.spi.open(0, 0)
+
+    self.spi.max_speed_hz = 30000000
+
+  # helpers
+  def _transfer(self, endpoint, data, max_rx_len=1000):
+    for _ in range(MAX_RETRY_COUNT):
+      try:
+        packet = struct.pack("<BBHH", SYNC, endpoint, len(data), max_rx_len)
+        packet += bytes([reduce(lambda x, y: x^y, packet) ^ CHECKSUM_START])
+        self.spi.xfer2(packet)
+
+        dat = b"\x00"
+        while dat[0] not in [HACK, NACK]:
+          dat = self.spi.xfer2(b"\x12")
+
+        if dat[0] == NACK:
+          raise Exception("Got NACK response for header")
+
+        packet = bytes(data)
+        packet += bytes([reduce(lambda x, y: x^y, packet) ^ CHECKSUM_START])
+        self.spi.xfer2(packet)
+
+        dat = b"\x00"
+        while dat[0] not in [DACK, NACK]:
+          dat = self.spi.xfer2(b"\xab")
+
+        if dat[0] == NACK:
+          raise Exception("Got NACK response for data")
+
+        response_len = struct.unpack("<H", bytes(self.spi.xfer2(b"\x00" * 2)))[0]
+
+        dat = bytes(self.spi.xfer2(b"\x00" * (response_len + 1)))
+        # TODO: verify CRC
+        dat = dat[:-1]
+
+        return dat
+      except Exception:
+        pass
+    raise Exception(f"SPI transaction failed {MAX_RETRY_COUNT} times")
+
+  # libusb1 functions
+  def close(self):
+    self.spi.close()
+
+  def controlWrite(self, request_type, request, value, index, data, timeout=0):
+    return self._transfer(0, struct.pack("<HHHH", request, value, index, 0))
+
+  def controlRead(self, request_type, request, value, index, length, timeout=0):
+    return self._transfer(0, struct.pack("<HHHH", request, value, index, length))
+
+  # TODO: implement these
+  def bulkWrite(self, endpoint, data, timeout=0):
+    pass
+
+  def bulkRead(self, endpoint, data, timeout=0):
+    pass
diff --git a/requirements.txt b/requirements.txt
index c7515326..f90aba96 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -14,3 +14,4 @@ pre-commit==2.13.0
 pylint==2.5.2
 scons
 flaky
+spidev
diff --git a/tests/health_test.py b/tests/health_test.py
index 86e56f94..1298d6b4 100755
--- a/tests/health_test.py
+++ b/tests/health_test.py
@@ -1,18 +1,23 @@
 #!/usr/bin/env python3
+import os
 import time
 from panda import Panda
 
 if __name__ == "__main__":
-  panda_serials = Panda.list()
-  pandas = []
-  for ps in panda_serials:
-    pandas.append(Panda(serial=ps))
-  if len(pandas) == 0:
-    print("No pandas connected")
-    assert False
+  spi = "SPI" in os.environ
+  print("using SPI" if spi else "using SPI")
 
+
+  i = 0
+  pi = 0
+
+  panda = Panda(spi=spi)
   while True:
-    for panda in pandas:
-      print(panda.health())
-    print("\n")
-    time.sleep(0.5)
+    st = time.monotonic()
+    while time.monotonic() - st < 1:
+      panda.health()
+      i += 1
+    print(i, panda.health(), "\n")
+    print(f"Speed: {i - pi}Hz")
+    pi = i
+