From a660e342fdcd42632af28bb3a08f0cf4137f5601 Mon Sep 17 00:00:00 2001
From: Adeeb Shihadeh <adeebshihadeh@gmail.com>
Date: Thu, 14 Sep 2023 12:49:59 -0700
Subject: [PATCH] erase log flash sectors in DFU recovery (#1661)

* erase log flash sectors in DFU recovery

* always erase by sector

* erase all

* fix
---
 python/base.py      | 7 +------
 python/constants.py | 4 ++++
 python/dfu.py       | 7 +++++--
 python/spi.py       | 6 ------
 python/usb.py       | 7 ++-----
 5 files changed, 12 insertions(+), 19 deletions(-)

diff --git a/python/base.py b/python/base.py
index 509d3ebf..ffe9353d 100644
--- a/python/base.py
+++ b/python/base.py
@@ -53,14 +53,9 @@ class BaseSTBootloaderHandle(ABC):
     ...
 
   @abstractmethod
-  def erase_app(self) -> None:
-    ...
-
-  @abstractmethod
-  def erase_bootstub(self) -> None:
+  def erase_sector(self, sector: int) -> None:
     ...
 
   @abstractmethod
   def jump(self, address: int) -> None:
     ...
-
diff --git a/python/constants.py b/python/constants.py
index 16409ac3..725b2e77 100644
--- a/python/constants.py
+++ b/python/constants.py
@@ -19,6 +19,10 @@ class McuConfig(NamedTuple):
   bootstub_address: int
   bootstub_fn: str
 
+  def sector_address(self, i):
+    # assume bootstub is in sector 0
+    return self.bootstub_address + sum(self.sector_sizes[:i])
+
 Fx = (
   0x1FFF7A10,
   0x800,
diff --git a/python/dfu.py b/python/dfu.py
index fc196e4a..661f9e8f 100644
--- a/python/dfu.py
+++ b/python/dfu.py
@@ -121,8 +121,11 @@ class PandaDFU:
 
   def program_bootstub(self, code_bootstub):
     self._handle.clear_status()
-    self._handle.erase_bootstub()
-    self._handle.erase_app()
+
+    # erase all sectors
+    for i in range(len(self._mcu_type.config.sector_sizes)):
+      self._handle.erase_sector(i)
+
     self._handle.program(self._mcu_type.config.bootstub_address, code_bootstub)
 
   def recover(self):
diff --git a/python/spi.py b/python/spi.py
index faec497c..988bcd05 100644
--- a/python/spi.py
+++ b/python/spi.py
@@ -403,12 +403,6 @@ class STBootloaderSPIHandle(BaseSTBootloaderHandle):
 
   # *** PandaDFU API ***
 
-  def erase_app(self):
-    self.erase_sector(1)
-
-  def erase_bootstub(self):
-    self.erase_sector(0)
-
   def get_mcu_type(self):
     return self._mcu_type
 
diff --git a/python/usb.py b/python/usb.py
index e6b807ca..a43bb51d 100644
--- a/python/usb.py
+++ b/python/usb.py
@@ -51,11 +51,8 @@ class STBootloaderUSBHandle(BaseSTBootloaderHandle):
   def get_mcu_type(self):
     return self._mcu_type
 
-  def erase_app(self):
-    self._erase_page_address(self._mcu_type.config.app_address)
-
-  def erase_bootstub(self):
-    self._erase_page_address(self._mcu_type.config.bootstub_address)
+  def erase_sector(self, sector: int):
+    self._erase_page_address(self._mcu_type.config.sector_address(sector))
 
   def clear_status(self):
     # Clear status