From 047d36398dc5a18e899a8120d1930d0e8eec8591 Mon Sep 17 00:00:00 2001
From: Adeeb Shihadeh <adeebshihadeh@gmail.com>
Date: Tue, 12 Sep 2023 20:04:54 -0700
Subject: [PATCH] SPI flash dump + restore tools (#1654)

* dump

* restore

* fix

* update

* update

* cleanup

---------

Co-authored-by: Comma Device <device@comma.ai>
---
 __init__.py                |  2 +-
 tests/read_flash_spi.py    | 32 ++++++++++++++++++++++++++++++++
 tests/restore_flash_spi.py | 27 +++++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 1 deletion(-)
 create mode 100755 tests/read_flash_spi.py
 create mode 100755 tests/restore_flash_spi.py

diff --git a/__init__.py b/__init__.py
index e8e05c52d..41876ee2a 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,5 +1,5 @@
 from .python.constants import McuType, BASEDIR, FW_PATH, USBPACKET_MAX_SIZE  # noqa: F401
-from .python.spi import PandaSpiException, PandaProtocolMismatch  # noqa: F401
+from .python.spi import PandaSpiException, PandaProtocolMismatch, STBootloaderSPIHandle  # noqa: F401
 from .python.serial import PandaSerial  # noqa: F401
 from .python.canhandle import CanHandle # noqa: F401
 from .python import (Panda, PandaDFU, # noqa: F401
diff --git a/tests/read_flash_spi.py b/tests/read_flash_spi.py
new file mode 100755
index 000000000..133062be7
--- /dev/null
+++ b/tests/read_flash_spi.py
@@ -0,0 +1,32 @@
+#!/usr/bin/env python3
+from panda import Panda, PandaDFU
+
+if __name__ == "__main__":
+  try:
+    from openpilot.system.hardware import HARDWARE
+    HARDWARE.recover_internal_panda()
+    Panda.wait_for_dfu(None, 5)
+  except Exception:
+    pass
+
+  p = PandaDFU(None)
+  cfg = p.get_mcu_type().config
+
+  def readmem(addr, length, fn):
+    print(f"reading {hex(addr)} {hex(length)} bytes to {fn}")
+    max_size = 255
+    with open(fn, "wb") as f:
+      to_read = length
+      while to_read > 0:
+        l = min(to_read, max_size)
+        dat = p._handle.read(addr, l)
+        assert len(dat) == l
+        f.write(dat)
+
+        to_read -= len(dat)
+        addr += len(dat)
+
+  addr = cfg.bootstub_address
+  for i, sector_size in enumerate(cfg.sector_sizes):
+    readmem(addr, sector_size, f"sector_{i}.bin")
+    addr += sector_size
diff --git a/tests/restore_flash_spi.py b/tests/restore_flash_spi.py
new file mode 100755
index 000000000..c23b29828
--- /dev/null
+++ b/tests/restore_flash_spi.py
@@ -0,0 +1,27 @@
+#!/usr/bin/env python3
+from panda import Panda, PandaDFU, STBootloaderSPIHandle
+
+if __name__ == "__main__":
+  try:
+    from openpilot.system.hardware import HARDWARE
+    HARDWARE.recover_internal_panda()
+    Panda.wait_for_dfu(None, 5)
+  except Exception:
+    pass
+
+  p = PandaDFU(None)
+  assert isinstance(p._handle, STBootloaderSPIHandle)
+  cfg = p.get_mcu_type().config
+
+  print("restoring from backup...")
+  addr = cfg.bootstub_address
+  for i, sector_size in enumerate(cfg.sector_sizes):
+    print(f"- sector #{i}")
+    p._handle.erase_sector(i)
+    with open(f"sector_{i}.bin", "rb") as f:
+      dat = f.read()
+      assert len(dat) == sector_size
+      p._handle.program(addr, dat)
+    addr += len(dat)
+
+  p.reset()