snagrecover ambarella

Author: phodina

src/snagrecover/config.py

@@ -27,6 +27,7 @@
 	"stm32mp": "0483:df11",
 	"sama5":    "03eb:6124",
 	"sunxi":    "1f3a:efe8",
+	"ambarella": "4255:0013",
 	"am6x":     "0451:6165",
 	"am62lx":   "0451:6165",
 	"zynqmp":   "03fd:0050",
@@ -114,4 +115,3 @@ def init_config(args: list):
 	# store input arguments in config
 	recovery_config["args"] = vars(args)
 	logger.debug(f"recovery_config:{str(recovery_config)}")
-

src/snagrecover/firmware/firmware.py

@@ -18,11 +18,33 @@
 # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 
 import logging
+from pathlib import Path
+from typing import Optional
+
 logger = logging.getLogger("snagrecover")
 from snagrecover.protocols import dfu
 from snagrecover.config import recovery_config
 from snagrecover.utils import cli_error
 
+class FirmwareError(Exception):
+    """Base class for firmware-related errors."""
+    pass
+
+class Firmware:
+    """Base class for firmware handlers."""
+    
+    def __init__(self):
+        """Initialize firmware handler."""
+        pass
+        
+    def load(self) -> None:
+        """Load firmware files.
+        
+        Raises:
+            FirmwareError: If loading fails
+        """
+        raise NotImplementedError("Subclasses must implement load()")
+
 def stm32mp_run(port, fw_name: str, fw_blob: bytes):
 	"""
 	There isn't a lot of complicated logic to handle stm32mp firmware

src/snagrecover/recoveries/ambarella.py

@@ -213,3 +213,41 @@ def flash_firmware(self, firmware: Path) -> None:
         with open(firmware, 'rb') as f:
             fw_data = f.read()
         self.protocol.send_file(fw_info.memfw_prog_addr, fw_data)
+
+def main() -> None:
+    """Main recovery function.
+    
+    Raises:
+        AmbaUsbError: On recovery error
+    """
+    from snagrecover.config import recovery_config
+    from snagrecover.utils import get_usb
+    
+    # Get USB device
+    usb_path = recovery_config["usb_path"]
+    device = UsbDevice(get_usb(usb_path))
+    
+    # Initialize recovery handler
+    recovery = AmbarellaRecovery(device)
+    
+    # Get firmware configuration
+    fw_config = recovery_config.get("firmware", {})
+    
+    # Initialize DRAM if script provided
+    if "dram_script" in fw_config:
+        dram_script = Path(fw_config["dram_script"]["path"])
+        recovery.initialize_dram(dram_script)
+    
+    # Load bootloader if provided
+    if "bootloader" in fw_config:
+        bootloader = Path(fw_config["bootloader"]["path"])
+        recovery.load_bootloader(bootloader)
+    
+    # Flash firmware if provided in steps
+    steps = recovery_config.get("steps", [])
+    for step in steps:
+        if step.get("action") == "flash_firmware":
+            firmware_path = step.get("args", {}).get("firmware")
+            if firmware_path:
+                firmware = Path(firmware_path)
+                recovery.flash_firmware(firmware)

src/snagrecover/usb.py

@@ -4,6 +4,57 @@
 import gc
 import importlib
 import sys
+from typing import Optional
+
+class UsbError(Exception):
+    """Base class for USB-related errors."""
+    pass
+
+class UsbDevice:
+    """USB device wrapper class."""
+    
+    def __init__(self, device: usb.core.Device):
+        """Initialize USB device.
+        
+        Args:
+            device: PyUSB device object
+        """
+        self.device = device
+        
+        # Set default configuration
+        try:
+            self.device.set_configuration()
+        except usb.core.USBError as e:
+            raise UsbError(f"Failed to set configuration: {e}")
+            
+    def write(self, data: bytes) -> None:
+        """Write data to device.
+        
+        Args:
+            data: Data to write
+            
+        Raises:
+            UsbError: On write error
+        """
+        try:
+            self.device.write(0x1, data)  # EP1 OUT
+        except usb.core.USBError as e:
+            raise UsbError(f"Write failed: {e}")
+            
+    def read(self, buffer: bytearray) -> None:
+        """Read data from device into buffer.
+        
+        Args:
+            buffer: Buffer to read into
+            
+        Raises:
+            UsbError: On read error
+        """
+        try:
+            read_data = self.device.read(0x81, len(buffer))  # EP1 IN
+            buffer[:len(read_data)] = read_data
+        except usb.core.USBError as e:
+            raise UsbError(f"Read failed: {e}")
 
 class SnagbootUSBContext():
 	"""
@@ -67,4 +118,3 @@ def find(**args):
 			tests = (hasattr(dev, key) and val == getattr(dev, key) for key, val in args.items())
 			if all(tests):
 				yield weakref.proxy(dev)
-