Fixed thumb mode bugs for reading/writing data

exploit.py

@@ -18,6 +18,10 @@ USBDEVFS_SUBMITURB = 0x8038550a
 USBDEVFS_REAPURB = 0x4008550c
 USBDEVFS_DISCARDURB = 0x0000550b
 
+IMEM_START = 0x40000000
+IMEM_SIZE = 0xfc00
+BOOTROM_START = 0x100000
+
 debug_exchanges = 0
 class DEVICE():
     def usb_connect(self):
@@ -115,6 +119,8 @@ class TegraRCM():
         if not IS_OSX:
             fds_before = get_fds()
         self.dev = DEVICE()
+        self.read = self.dev.read
+        self.write = self.dev.write
         if not IS_OSX:
             self.fds = get_fds() - fds_before
             self.fd = sorted(list(self.fds))[-1]
@@ -161,7 +167,10 @@ class TegraRCM():
     def build_rcm_cmd(self, payload_file_fd,  rcm_cmd_buf, rcm_cmd_buf_len, payload_thumb_mode ):
         ret = -1
         rcm_cmd_len = struct.pack("<L", RCM_CMD_LEN)
-        payload_entry = struct.pack("<L", (BOOTROM_PAYLOAD_ENTRY | 0x1))
+        if payload_thumb_mode:
+            payload_entry = struct.pack("<L", (BOOTROM_PAYLOAD_ENTRY | 0x1))
+        else:
+            payload_entry = struct.pack("<L", (BOOTROM_PAYLOAD_ENTRY))
         payload_thumb_mode = struct.pack("<L", payload_thumb_mode)
         self.read_intermezzo(rcm_cmd_buf)
         self.read_payload_file(payload_file_fd, rcm_cmd_buf, rcm_cmd_buf_len)
@@ -206,39 +215,42 @@ class TegraRCM():
         if(status == 0):
             error("wrong status returned!")
 
-    def send_verify_cmd(self, cmd):
-        self.dev.write(cmd)
-        r = self.dev.read(0x200)
-        if(r != cmd):
-            error(f"Error on sending command! {r}")
-            return False
-        return True
-
-    def handle_done(self):
-        r = self.dev.read(0x200)
-        if(r != b"done"):
-            error("Error on writing vbar!")
-
     def memdump_region(self, offset, size):
-        if(not self.send_verify_cmd(b"PEEK")):
-            return
-        mem_param = struct.pack('<LL', offset, size)
-        self.dev.write(mem_param)
+        '''
+        Dump a region from target device. Based on an offset/address and size:
+
+        Args:
+
+            :param offset: Address of which to dump
+            :param size: Size to dump
+
+        Returns:
+
+            Bytes
+        '''
+        self.write(b"PEEK")
+        mem_param = struct.pack('<III', offset, size, 0) #Send extra 4 bytes to fill the 12 byte buffer
+        self.write(mem_param)
         received = b''
-        blk_sz = 0x200
+        blk_sz = 0x100
         while len(received) < size:
-            if (remaining := size - len(received)) < 0x200:
+            if (remaining := size - len(received)) < blk_sz:
                 blk_sz = remaining
-            d = self.dev.read(blk_sz)
+            d = self.read(blk_sz)
             if len(d) == blk_sz:
-                self.dev.write(b"ACK\x00")
+                self.write(b"ACK\x00")
             received += d
-        self.handle_done()
+        while True:
+            try:
+                self.write(b"ACK\x00")
+                break
+            except:
+                pass
         return received
     
-    def memwrite_region(self, address, data, check=True):
+    def memwrite_region(self, address, data):
         '''
-        Write a blob of data to an address on the device. Sometimes this function has issues when writing more than 0x20 bytes of data
+        Write a blob of data to an address on the device.
 
         Args:
 
@@ -247,32 +259,23 @@ class TegraRCM():
             :param (Bool): check if data is really written by dumping the region and checking if it has changed
         '''
         size = len(data)
-        if(check):
-            before = self.memdump_region(address, size)
-
-        if(not self.send_verify_cmd(b"POKE")):
-            return
-        mem_param = struct.pack('<II', address, size)
-        self.dev.write(mem_param)
+        self.write(b"POKE")
+        mem_param = struct.pack('<III', address, size, 0) #Send extra 4 bytes to fill the 12 byte buffer
+        self.write(mem_param)
+        blk_sz = 0x200
 
         while len(data) > 0:
-            remaining = 0x200
-            if(len(data) < 0x200):
+            remaining = blk_sz
+            if(len(data) < blk_sz):
                 remaining = len(data)
             send = data[:remaining]
             data = data[remaining:]
-            self.dev.write(send)
-            message = self.dev.read(0x200)
+            self.write(send)
+            message = self.read(blk_sz)
             if(message != b"OK"):
                 error("Error on writing data to device!")
                 return
-            self.dev.write(b"ACK\x00")
-        self.handle_done()
-        #Read back data
-        if(check):
-            after = self.memdump_region(address, size)
-            if(after == before and send != before):
-                error(f"Memory written succesfully, but no changes detected! | {hex(address)}")
+            self.write(b"ACK\x00")
 
     def search_bootrom(self):
         dumped = BytesIO()
@@ -297,7 +300,7 @@ class TegraRCM():
                 self.search_bootrom()
                 #dump memory
                 self.dump_bootrom()
-            
+ 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -318,5 +321,9 @@ if __name__ == "__main__":
         # d2 = rcm.dev.read(0x200)
         if d == b"GiAs": 
             ok("Device in GA debugger")
+            data = rcm.memdump_region(IMEM_START, 0x1000)
+            data = rcm.memdump_region(IMEM_START, 0x1000)
+            rcm.memwrite_region(IMEM_START, b"\xaf" * 0x100)
+            data2 = rcm.memdump_region(IMEM_START, 0x1000)
     else:
         rcm.cmd_handler()