diff --git a/ubireader/scripts/ubireader_display_blocks.py b/ubireader/scripts/ubireader_display_blocks.py index 8273219..eec5785 100755 --- a/ubireader/scripts/ubireader_display_blocks.py +++ b/ubireader/scripts/ubireader_display_blocks.py @@ -109,19 +109,21 @@ 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.') @@ -129,9 +131,9 @@ 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.') @@ -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)) diff --git a/ubireader/scripts/ubireader_display_info.py b/ubireader/scripts/ubireader_display_info.py index b5291aa..889566f 100755 --- a/ubireader/scripts/ubireader_display_info.py +++ b/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__': diff --git a/ubireader/scripts/ubireader_extract_files.py b/ubireader/scripts/ubireader_extract_files.py index a18a214..59d2ebe 100755 --- a/ubireader/scripts/ubireader_extract_files.py +++ b/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,9 +136,9 @@ 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.') @@ -144,7 +146,7 @@ def main(): 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__': diff --git a/ubireader/scripts/ubireader_extract_images.py b/ubireader/scripts/ubireader_extract_images.py index 3d8cf2e..dfb854c 100755 --- a/ubireader/scripts/ubireader_extract_images.py +++ b/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__': diff --git a/ubireader/scripts/ubireader_list_files.py b/ubireader/scripts/ubireader_list_files.py index 1ac00be..000498f 100755 --- a/ubireader/scripts/ubireader_list_files.py +++ b/ubireader/scripts/ubireader_list_files.py @@ -101,19 +101,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.') @@ -121,15 +123,15 @@ 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 @@ -174,6 +176,7 @@ def main(): print('Something went wrong to get here.') ufile_obj.close() + fileobj.close() if __name__=='__main__': diff --git a/ubireader/scripts/ubireader_utils_info.py b/ubireader/scripts/ubireader_utils_info.py index b8fafe5..2dc5afb 100755 --- a/ubireader/scripts/ubireader_utils_info.py +++ b/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__': diff --git a/ubireader/ubi_io.py b/ubireader/ubi_io.py index d6f80db..10d4deb 100755 --- a/ubireader/ubi_io.py +++ b/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) diff --git a/ubireader/utils.py b/ubireader/utils.py index 1793285..c479daf 100755 --- a/ubireader/utils.py +++ b/ubireader/utils.py @@ -26,7 +26,8 @@ def guess_start_offset(path, guess_offset=0): file_offset = guess_offset - f = open(path, 'rb') + is_fileobj = hasattr(path, "read") + f = path if is_fileobj else open(path, 'rb') f.seek(0,2) file_size = f.tell()+1 f.seek(guess_offset) @@ -57,26 +58,31 @@ def guess_start_offset(path, guess_offset=0): else: error(guess_start_offset, 'Fatal', 'Could not determine start offset.') - f.close() + if not is_fileobj: + f.close() def guess_filetype(path, start_offset=0): log(guess_filetype, 'Looking for file type at %s' % start_offset) - with open(path, 'rb') as f: - f.seek(start_offset) - buf = f.read(4) + is_fileobj = hasattr(path, "read") + f = path if is_fileobj else open(path, 'rb') + f.seek(start_offset) + buf = f.read(4) - if buf == UBI_EC_HDR_MAGIC: - ftype = UBI_EC_HDR_MAGIC - log(guess_filetype, 'File looks like a UBI image.') + if buf == UBI_EC_HDR_MAGIC: + ftype = UBI_EC_HDR_MAGIC + log(guess_filetype, 'File looks like a UBI image.') - elif buf == UBIFS_NODE_MAGIC: - ftype = UBIFS_NODE_MAGIC - log(guess_filetype, 'File looks like a UBIFS image.') - else: - ftype = None - error(guess_filetype, 'Fatal', 'Could not determine file type.') + elif buf == UBIFS_NODE_MAGIC: + ftype = UBIFS_NODE_MAGIC + log(guess_filetype, 'File looks like a UBIFS image.') + else: + ftype = None + error(guess_filetype, 'Fatal', 'Could not determine file type.') + + if not is_fileobj: + f.close() return ftype @@ -93,7 +99,8 @@ def guess_leb_size(path): Searches file for superblock and retrieves leb size. """ - f = open(path, 'rb') + is_fileobj = hasattr(path, "read") + f = path if is_fileobj else open(path, 'rb') f.seek(0,2) file_size = f.tell()+1 f.seek(0) @@ -118,10 +125,12 @@ def guess_leb_size(path): sbn = nodes.sb_node(buf) block_size = sbn.leb_size - f.close() + if not is_fileobj: + f.close() return block_size - f.close() + if not is_fileobj: + f.close() return block_size @@ -139,7 +148,8 @@ def guess_peb_size(path): """ file_offset = 0 offsets = [] - f = open(path, 'rb') + is_fileobj = hasattr(path, "read") + f = path if is_fileobj else open(path, 'rb') f.seek(0,2) file_size = f.tell()+1 f.seek(0) @@ -158,7 +168,8 @@ def guess_peb_size(path): offsets.append(idx) file_offset += FILE_CHUNK_SZ - f.close() + if not is_fileobj: + f.close() occurrences = {} for i in range(0, len(offsets)):