From 02773d19acafab21735f381e13e9b72c2017166b Mon Sep 17 00:00:00 2001 From: AT0myks Date: Tue, 3 Dec 2024 16:04:59 +0000 Subject: [PATCH] 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 | 15 +++--- ubireader/scripts/ubireader_display_info.py | 15 +++--- ubireader/scripts/ubireader_extract_files.py | 15 +++--- ubireader/scripts/ubireader_extract_images.py | 13 +++-- ubireader/scripts/ubireader_list_files.py | 15 +++--- ubireader/scripts/ubireader_utils_info.py | 13 +++-- ubireader/ubi_io.py | 17 ++++--- ubireader/utils.py | 49 ++++++++++++------- 8 files changed, 93 insertions(+), 59 deletions(-) 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)):