Accept file objects

For functions that work on UBI files and call
`open()`, additionally allow a file object to be
passed, and in that case leave the responsibility
of closing it to the caller. This means that the
scripts can now open the file once instead of two
to four times previously.

ubireader/scripts/ubireader_display_blocks.py

@@ -109,29 +109,31 @@ def main():
         parser.error('File path must be provided.')
         sys.exit(1)
 
+    fileobj = open(path, "rb")
+
     if args.start_offset:
         start_offset = args.start_offset
     elif args.guess_offset:
-        start_offset = guess_start_offset(path, args.guess_offset)
+        start_offset = guess_start_offset(fileobj, args.guess_offset)
     else:
-        start_offset = guess_start_offset(path)
+        start_offset = guess_start_offset(fileobj)
 
     if args.end_offset:
         end_offset = args.end_offset
     else:
         end_offset = None
 
-    filetype = guess_filetype(path, start_offset)
+    filetype = guess_filetype(fileobj, start_offset)
     if not filetype:
         parser.error('Could not determine file type.')
 
     if args.block_size:
         block_size = args.block_size
     else:
         if filetype == UBI_EC_HDR_MAGIC:
-            block_size = guess_peb_size(path)
+            block_size = guess_peb_size(fileobj)
         elif filetype == UBIFS_NODE_MAGIC:
-            block_size = guess_leb_size(path)
+            block_size = guess_leb_size(fileobj)
 
         if not block_size:
             parser.error('Block size could not be determined.')
@@ -153,7 +155,7 @@ def main():
         parser.error('No search parameters given, -b arg is required.')
 
 
-    ufile_obj = ubi_file(path, block_size, start_offset, end_offset)
+    ufile_obj = ubi_file(fileobj, block_size, start_offset, end_offset)
     ubi_obj = ubi_base(ufile_obj)
     blocks = []
 
@@ -185,6 +187,7 @@ def main():
             blocks.append(ubi_obj.blocks[block])
 
     ufile_obj.close()
+    fileobj.close()
 
     print('\nBlock matches: %s' % len(blocks))
 

ubireader/scripts/ubireader_display_info.py

@@ -93,19 +93,21 @@ def main():
         if not os.path.exists(path):
             parser.error("File path doesn't exist.")
 
+    fileobj = open(path, "rb")
+
     if args.start_offset:
         start_offset = args.start_offset
     elif args.guess_offset:
-        start_offset = guess_start_offset(path, args.guess_offset)
+        start_offset = guess_start_offset(fileobj, args.guess_offset)
     else:
-        start_offset = guess_start_offset(path)
+        start_offset = guess_start_offset(fileobj)
 
     if args.end_offset:
         end_offset = args.end_offset
     else:
         end_offset = None
 
-    filetype = guess_filetype(path, start_offset)
+    filetype = guess_filetype(fileobj, start_offset)
     if not filetype:
         parser.error('Could not determine file type.')
 
@@ -115,16 +117,16 @@ def main():
         block_size = args.block_size
     else:
         if filetype == UBI_EC_HDR_MAGIC:
-            block_size = guess_peb_size(path)
+            block_size = guess_peb_size(fileobj)
         elif filetype == UBIFS_NODE_MAGIC:
-            block_size = guess_leb_size(path)
+            block_size = guess_leb_size(fileobj)
 
         if not block_size:
             parser.error('Block size could not be determined.')
 
 
     # Create file object.
-    ufile_obj = ubi_file(path, block_size, start_offset, end_offset)
+    ufile_obj = ubi_file(fileobj, block_size, start_offset, end_offset)
 
     if filetype == UBI_EC_HDR_MAGIC:
         # Create UBI object
@@ -186,6 +188,7 @@ def main():
         print('Something went wrong to get here.')
 
     ufile_obj.close()
+    fileobj.close()
 
 
 if __name__=='__main__':

ubireader/scripts/ubireader_extract_files.py

@@ -109,19 +109,21 @@ def main():
         if not os.path.exists(path):
             parser.error("File path doesn't exist.")
 
+    fileobj = open(path, "rb")
+
     if args.start_offset:
         start_offset = args.start_offset
     elif args.guess_offset:
-        start_offset = guess_start_offset(path, args.guess_offset)
+        start_offset = guess_start_offset(fileobj, args.guess_offset)
     else:
-        start_offset = guess_start_offset(path)
+        start_offset = guess_start_offset(fileobj)
 
     if args.end_offset:
         end_offset = args.end_offset
     else:
         end_offset = None
 
-    filetype = guess_filetype(path, start_offset)
+    filetype = guess_filetype(fileobj, start_offset)
     if not filetype:
         parser.error('Could not determine file type.')
 
@@ -134,17 +136,17 @@ def main():
         block_size = args.block_size
     else:
         if filetype == UBI_EC_HDR_MAGIC:
-            block_size = guess_peb_size(path)
+            block_size = guess_peb_size(fileobj)
         elif filetype == UBIFS_NODE_MAGIC:
-            block_size = guess_leb_size(path)
+            block_size = guess_leb_size(fileobj)
 
         if not block_size:
             parser.error('Block size could not be determined.')
 
     perms = args.permissions
 
     # Create file object.
-    ufile_obj = ubi_file(path, block_size, start_offset, end_offset)
+    ufile_obj = ubi_file(fileobj, block_size, start_offset, end_offset)
 
     if filetype == UBI_EC_HDR_MAGIC:
         # Create UBI object
@@ -199,6 +201,7 @@ def main():
         print('Something went wrong to get here.')
 
     ufile_obj.close()
+    fileobj.close()
 
 
 if __name__=='__main__':

ubireader/scripts/ubireader_extract_images.py

@@ -100,19 +100,21 @@ def main():
         if not os.path.exists(path):
             parser.error("File path doesn't exist.")
 
+    fileobj = open(path, "rb")
+
     if args.start_offset:
         start_offset = args.start_offset
     elif args.guess_offset:
-        start_offset = guess_start_offset(path, args.guess_offset)
+        start_offset = guess_start_offset(fileobj, args.guess_offset)
     else:
-        start_offset = guess_start_offset(path)
+        start_offset = guess_start_offset(fileobj)
 
     if args.end_offset:
         end_offset = args.end_offset
     else:
         end_offset = None
 
-    filetype = guess_filetype(path, start_offset)
+    filetype = guess_filetype(fileobj, start_offset)
     if filetype != UBI_EC_HDR_MAGIC:
         parser.error('File does not look like UBI data.')
 
@@ -125,7 +127,7 @@ def main():
     if args.block_size:
         block_size = args.block_size
     else:
-        block_size = guess_peb_size(path)
+        block_size = guess_peb_size(fileobj)
 
         if not block_size:
             parser.error('Block size could not be determined.')
@@ -136,7 +138,7 @@ def main():
         image_type = 'UBIFS'
 
     # Create file object.
-    ufile_obj = ubi_file(path, block_size, start_offset, end_offset)
+    ufile_obj = ubi_file(fileobj, block_size, start_offset, end_offset)
 
     # Create UBI object
     ubi_obj = ubi(ufile_obj)
@@ -168,6 +170,7 @@ def main():
                     f.write(block)
 
     ufile_obj.close()
+    fileobj.close()
 
 
 if __name__=='__main__':

ubireader/scripts/ubireader_list_files.py

@@ -101,35 +101,37 @@ def main():
         if not os.path.exists(path):
             parser.error("File path doesn't exist.")
 
+    fileobj = open(path, "rb")
+
     if args.start_offset:
         start_offset = args.start_offset
     elif args.guess_offset:
-        start_offset = guess_start_offset(path, args.guess_offset)
+        start_offset = guess_start_offset(fileobj, args.guess_offset)
     else:
-        start_offset = guess_start_offset(path)
+        start_offset = guess_start_offset(fileobj)
 
     if args.end_offset:
         end_offset = args.end_offset
     else:
         end_offset = None
 
-    filetype = guess_filetype(path, start_offset)
+    filetype = guess_filetype(fileobj, start_offset)
     if not filetype:
         parser.error('Could not determine file type.')
 
     if args.block_size:
         block_size = args.block_size
     else:
         if filetype == UBI_EC_HDR_MAGIC:
-            block_size = guess_peb_size(path)
+            block_size = guess_peb_size(fileobj)
         elif filetype == UBIFS_NODE_MAGIC:
-            block_size = guess_leb_size(path)
+            block_size = guess_leb_size(fileobj)
 
         if not block_size:
             parser.error('Block size could not be determined.')
 
     # Create file object.
-    ufile_obj = ubi_file(path, block_size, start_offset, end_offset)
+    ufile_obj = ubi_file(fileobj, block_size, start_offset, end_offset)
 
     if filetype == UBI_EC_HDR_MAGIC:
         # Create UBI object
@@ -174,6 +176,7 @@ def main():
         print('Something went wrong to get here.')
 
     ufile_obj.close()
+    fileobj.close()
 
 
 if __name__=='__main__':

ubireader/scripts/ubireader_utils_info.py

@@ -299,19 +299,21 @@ def main():
         if not os.path.exists(path):
             parser.error("File path doesn't exist.")
 
+    fileobj = open(path, "rb")
+
     if args.start_offset:
         start_offset = args.start_offset
     elif args.guess_offset:
-        start_offset = guess_start_offset(path, args.guess_offset)
+        start_offset = guess_start_offset(fileobj, args.guess_offset)
     else:
-        start_offset = guess_start_offset(path)
+        start_offset = guess_start_offset(fileobj)
 
     if args.end_offset:
         end_offset = args.end_offset
     else:
         end_offset = None
 
-    filetype = guess_filetype(path, start_offset)
+    filetype = guess_filetype(fileobj, start_offset)
     if filetype != UBI_EC_HDR_MAGIC:
         parser.error('File does not look like UBI data.')
 
@@ -324,13 +326,13 @@ def main():
     if args.block_size:
         block_size = args.block_size
     else:
-        block_size = guess_peb_size(path)
+        block_size = guess_peb_size(fileobj)
 
         if not block_size:
             parser.error('Block size could not be determined.')
 
     # Create file object.
-    ufile_obj = ubi_file(path, block_size, start_offset, end_offset)
+    ufile_obj = ubi_file(fileobj, block_size, start_offset, end_offset)
 
     # Create UBI object
     ubi_obj = ubi(ufile_obj)
@@ -344,6 +346,7 @@ def main():
         make_files(ubi_obj, outpath)
 
     ufile_obj.close()
+    fileobj.close()
 
 
 if __name__=='__main__':

ubireader/ubi_io.py

@@ -52,11 +52,15 @@ class ubi_file(object):
     def __init__(self, path, block_size, start_offset=0, end_offset=None):
         self.__name__ = 'UBI_File'
         self.is_valid = False
-        try:
-            log(self, 'Open Path: %s' % path)
-            self._fhandle = open(path, 'rb')
-        except Exception as e:
-            error(self, 'Fatal', 'Open file: %s' % e)
+        self._fileobj_passed = hasattr(path, "read")
+        if self._fileobj_passed:
+            self._fhandle = path
+        else:
+            try:
+                log(self, 'Open Path: %s' % path)
+                self._fhandle = open(path, 'rb')
+            except Exception as e:
+                error(self, 'Fatal', 'Open file: %s' % e)
 
         self._fhandle.seek(0,2)
         file_size = self.tell()
@@ -111,7 +115,8 @@ def _get_block_size(self):
     block_size = property(_get_block_size)
 
     def close(self):
-        self._fhandle.close()
+        if not self._fileobj_passed:
+            self._fhandle.close()
 
     def seek(self, offset):
         self._fhandle.seek(offset)