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 a5bc154..46049b6 100755 --- a/ubireader/scripts/ubireader_extract_files.py +++ b/ubireader/scripts/ubireader_extract_files.py @@ -125,19 +125,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.') @@ -150,9 +152,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.') @@ -160,7 +162,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 @@ -215,6 +217,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 041e33a..1151b46 100755 --- a/ubireader/scripts/ubireader_list_files.py +++ b/ubireader/scripts/ubireader_list_files.py @@ -142,19 +142,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.') @@ -162,15 +164,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 @@ -215,6 +217,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 fd49621..c4c8aff 100755 --- a/ubireader/ubi_io.py +++ b/ubireader/ubi_io.py @@ -61,11 +61,15 @@ class ubi_file(object): def __init__(self, path: str, block_size: int, start_offset: int = 0, end_offset: int | None = None) -> 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() @@ -120,7 +124,8 @@ def _get_block_size(self) -> int: block_size = property(_get_block_size) def close(self) -> None: - self._fhandle.close() + if not self._fileobj_passed: + self._fhandle.close() def seek(self, offset: int) -> None: self._fhandle.seek(offset) diff --git a/ubireader/utils.py b/ubireader/utils.py index ce57464..705a89a 100755 --- a/ubireader/utils.py +++ b/ubireader/utils.py @@ -27,7 +27,8 @@ def guess_start_offset(path: str, guess_offset: int =0) -> int | None: 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) @@ -58,26 +59,31 @@ def guess_start_offset(path: str, guess_offset: int =0) -> int | None: else: error(guess_start_offset, 'Fatal', 'Could not determine start offset.') - f.close() + if not is_fileobj: + f.close() def guess_filetype(path: str, start_offset: int = 0) -> bytes | None: 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 @@ -94,7 +100,8 @@ def guess_leb_size(path: str) -> int | None: 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) @@ -119,10 +126,12 @@ def guess_leb_size(path: str) -> int | None: 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 @@ -140,7 +149,8 @@ def guess_peb_size(path: str) -> int | None: """ file_offset = 0 offsets: list[int] = [] - 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) @@ -159,7 +169,8 @@ def guess_peb_size(path: str) -> int | None: offsets.append(idx) file_offset += FILE_CHUNK_SZ - f.close() + if not is_fileobj: + f.close() occurrences: dict[int, int] = {} for i in range(0, len(offsets)):