Skip to content

Commit 02773d1

Browse files
AT0myksAT0myks
committed
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.
1 parent cc7225e commit 02773d1

8 files changed

+93
-59
lines changed

‎ubireader/scripts/ubireader_display_blocks.py

+9-6
Original file line numberDiff line numberDiff line change
@@ -109,29 +109,31 @@ def main():
109109
parser.error('File path must be provided.')
110110
sys.exit(1)
111111

112+
fileobj = open(path, "rb")
113+
112114
if args.start_offset:
113115
start_offset = args.start_offset
114116
elif args.guess_offset:
115-
start_offset = guess_start_offset(path, args.guess_offset)
117+
start_offset = guess_start_offset(fileobj, args.guess_offset)
116118
else:
117-
start_offset = guess_start_offset(path)
119+
start_offset = guess_start_offset(fileobj)
118120

119121
if args.end_offset:
120122
end_offset = args.end_offset
121123
else:
122124
end_offset = None
123125

124-
filetype = guess_filetype(path, start_offset)
126+
filetype = guess_filetype(fileobj, start_offset)
125127
if not filetype:
126128
parser.error('Could not determine file type.')
127129

128130
if args.block_size:
129131
block_size = args.block_size
130132
else:
131133
if filetype == UBI_EC_HDR_MAGIC:
132-
block_size = guess_peb_size(path)
134+
block_size = guess_peb_size(fileobj)
133135
elif filetype == UBIFS_NODE_MAGIC:
134-
block_size = guess_leb_size(path)
136+
block_size = guess_leb_size(fileobj)
135137

136138
if not block_size:
137139
parser.error('Block size could not be determined.')
@@ -153,7 +155,7 @@ def main():
153155
parser.error('No search parameters given, -b arg is required.')
154156

155157

156-
ufile_obj = ubi_file(path, block_size, start_offset, end_offset)
158+
ufile_obj = ubi_file(fileobj, block_size, start_offset, end_offset)
157159
ubi_obj = ubi_base(ufile_obj)
158160
blocks = []
159161

@@ -185,6 +187,7 @@ def main():
185187
blocks.append(ubi_obj.blocks[block])
186188

187189
ufile_obj.close()
190+
fileobj.close()
188191

189192
print('\nBlock matches: %s' % len(blocks))
190193

‎ubireader/scripts/ubireader_display_info.py

+9-6
Original file line numberDiff line numberDiff line change
@@ -93,19 +93,21 @@ def main():
9393
if not os.path.exists(path):
9494
parser.error("File path doesn't exist.")
9595

96+
fileobj = open(path, "rb")
97+
9698
if args.start_offset:
9799
start_offset = args.start_offset
98100
elif args.guess_offset:
99-
start_offset = guess_start_offset(path, args.guess_offset)
101+
start_offset = guess_start_offset(fileobj, args.guess_offset)
100102
else:
101-
start_offset = guess_start_offset(path)
103+
start_offset = guess_start_offset(fileobj)
102104

103105
if args.end_offset:
104106
end_offset = args.end_offset
105107
else:
106108
end_offset = None
107109

108-
filetype = guess_filetype(path, start_offset)
110+
filetype = guess_filetype(fileobj, start_offset)
109111
if not filetype:
110112
parser.error('Could not determine file type.')
111113

@@ -115,16 +117,16 @@ def main():
115117
block_size = args.block_size
116118
else:
117119
if filetype == UBI_EC_HDR_MAGIC:
118-
block_size = guess_peb_size(path)
120+
block_size = guess_peb_size(fileobj)
119121
elif filetype == UBIFS_NODE_MAGIC:
120-
block_size = guess_leb_size(path)
122+
block_size = guess_leb_size(fileobj)
121123

122124
if not block_size:
123125
parser.error('Block size could not be determined.')
124126

125127

126128
# Create file object.
127-
ufile_obj = ubi_file(path, block_size, start_offset, end_offset)
129+
ufile_obj = ubi_file(fileobj, block_size, start_offset, end_offset)
128130

129131
if filetype == UBI_EC_HDR_MAGIC:
130132
# Create UBI object
@@ -186,6 +188,7 @@ def main():
186188
print('Something went wrong to get here.')
187189

188190
ufile_obj.close()
191+
fileobj.close()
189192

190193

191194
if __name__=='__main__':

‎ubireader/scripts/ubireader_extract_files.py

+9-6
Original file line numberDiff line numberDiff line change
@@ -109,19 +109,21 @@ def main():
109109
if not os.path.exists(path):
110110
parser.error("File path doesn't exist.")
111111

112+
fileobj = open(path, "rb")
113+
112114
if args.start_offset:
113115
start_offset = args.start_offset
114116
elif args.guess_offset:
115-
start_offset = guess_start_offset(path, args.guess_offset)
117+
start_offset = guess_start_offset(fileobj, args.guess_offset)
116118
else:
117-
start_offset = guess_start_offset(path)
119+
start_offset = guess_start_offset(fileobj)
118120

119121
if args.end_offset:
120122
end_offset = args.end_offset
121123
else:
122124
end_offset = None
123125

124-
filetype = guess_filetype(path, start_offset)
126+
filetype = guess_filetype(fileobj, start_offset)
125127
if not filetype:
126128
parser.error('Could not determine file type.')
127129

@@ -134,17 +136,17 @@ def main():
134136
block_size = args.block_size
135137
else:
136138
if filetype == UBI_EC_HDR_MAGIC:
137-
block_size = guess_peb_size(path)
139+
block_size = guess_peb_size(fileobj)
138140
elif filetype == UBIFS_NODE_MAGIC:
139-
block_size = guess_leb_size(path)
141+
block_size = guess_leb_size(fileobj)
140142

141143
if not block_size:
142144
parser.error('Block size could not be determined.')
143145

144146
perms = args.permissions
145147

146148
# Create file object.
147-
ufile_obj = ubi_file(path, block_size, start_offset, end_offset)
149+
ufile_obj = ubi_file(fileobj, block_size, start_offset, end_offset)
148150

149151
if filetype == UBI_EC_HDR_MAGIC:
150152
# Create UBI object
@@ -199,6 +201,7 @@ def main():
199201
print('Something went wrong to get here.')
200202

201203
ufile_obj.close()
204+
fileobj.close()
202205

203206

204207
if __name__=='__main__':

‎ubireader/scripts/ubireader_extract_images.py

+8-5
Original file line numberDiff line numberDiff line change
@@ -100,19 +100,21 @@ def main():
100100
if not os.path.exists(path):
101101
parser.error("File path doesn't exist.")
102102

103+
fileobj = open(path, "rb")
104+
103105
if args.start_offset:
104106
start_offset = args.start_offset
105107
elif args.guess_offset:
106-
start_offset = guess_start_offset(path, args.guess_offset)
108+
start_offset = guess_start_offset(fileobj, args.guess_offset)
107109
else:
108-
start_offset = guess_start_offset(path)
110+
start_offset = guess_start_offset(fileobj)
109111

110112
if args.end_offset:
111113
end_offset = args.end_offset
112114
else:
113115
end_offset = None
114116

115-
filetype = guess_filetype(path, start_offset)
117+
filetype = guess_filetype(fileobj, start_offset)
116118
if filetype != UBI_EC_HDR_MAGIC:
117119
parser.error('File does not look like UBI data.')
118120

@@ -125,7 +127,7 @@ def main():
125127
if args.block_size:
126128
block_size = args.block_size
127129
else:
128-
block_size = guess_peb_size(path)
130+
block_size = guess_peb_size(fileobj)
129131

130132
if not block_size:
131133
parser.error('Block size could not be determined.')
@@ -136,7 +138,7 @@ def main():
136138
image_type = 'UBIFS'
137139

138140
# Create file object.
139-
ufile_obj = ubi_file(path, block_size, start_offset, end_offset)
141+
ufile_obj = ubi_file(fileobj, block_size, start_offset, end_offset)
140142

141143
# Create UBI object
142144
ubi_obj = ubi(ufile_obj)
@@ -168,6 +170,7 @@ def main():
168170
f.write(block)
169171

170172
ufile_obj.close()
173+
fileobj.close()
171174

172175

173176
if __name__=='__main__':

‎ubireader/scripts/ubireader_list_files.py

+9-6
Original file line numberDiff line numberDiff line change
@@ -101,35 +101,37 @@ def main():
101101
if not os.path.exists(path):
102102
parser.error("File path doesn't exist.")
103103

104+
fileobj = open(path, "rb")
105+
104106
if args.start_offset:
105107
start_offset = args.start_offset
106108
elif args.guess_offset:
107-
start_offset = guess_start_offset(path, args.guess_offset)
109+
start_offset = guess_start_offset(fileobj, args.guess_offset)
108110
else:
109-
start_offset = guess_start_offset(path)
111+
start_offset = guess_start_offset(fileobj)
110112

111113
if args.end_offset:
112114
end_offset = args.end_offset
113115
else:
114116
end_offset = None
115117

116-
filetype = guess_filetype(path, start_offset)
118+
filetype = guess_filetype(fileobj, start_offset)
117119
if not filetype:
118120
parser.error('Could not determine file type.')
119121

120122
if args.block_size:
121123
block_size = args.block_size
122124
else:
123125
if filetype == UBI_EC_HDR_MAGIC:
124-
block_size = guess_peb_size(path)
126+
block_size = guess_peb_size(fileobj)
125127
elif filetype == UBIFS_NODE_MAGIC:
126-
block_size = guess_leb_size(path)
128+
block_size = guess_leb_size(fileobj)
127129

128130
if not block_size:
129131
parser.error('Block size could not be determined.')
130132

131133
# Create file object.
132-
ufile_obj = ubi_file(path, block_size, start_offset, end_offset)
134+
ufile_obj = ubi_file(fileobj, block_size, start_offset, end_offset)
133135

134136
if filetype == UBI_EC_HDR_MAGIC:
135137
# Create UBI object
@@ -174,6 +176,7 @@ def main():
174176
print('Something went wrong to get here.')
175177

176178
ufile_obj.close()
179+
fileobj.close()
177180

178181

179182
if __name__=='__main__':

‎ubireader/scripts/ubireader_utils_info.py

+8-5
Original file line numberDiff line numberDiff line change
@@ -299,19 +299,21 @@ def main():
299299
if not os.path.exists(path):
300300
parser.error("File path doesn't exist.")
301301

302+
fileobj = open(path, "rb")
303+
302304
if args.start_offset:
303305
start_offset = args.start_offset
304306
elif args.guess_offset:
305-
start_offset = guess_start_offset(path, args.guess_offset)
307+
start_offset = guess_start_offset(fileobj, args.guess_offset)
306308
else:
307-
start_offset = guess_start_offset(path)
309+
start_offset = guess_start_offset(fileobj)
308310

309311
if args.end_offset:
310312
end_offset = args.end_offset
311313
else:
312314
end_offset = None
313315

314-
filetype = guess_filetype(path, start_offset)
316+
filetype = guess_filetype(fileobj, start_offset)
315317
if filetype != UBI_EC_HDR_MAGIC:
316318
parser.error('File does not look like UBI data.')
317319

@@ -324,13 +326,13 @@ def main():
324326
if args.block_size:
325327
block_size = args.block_size
326328
else:
327-
block_size = guess_peb_size(path)
329+
block_size = guess_peb_size(fileobj)
328330

329331
if not block_size:
330332
parser.error('Block size could not be determined.')
331333

332334
# Create file object.
333-
ufile_obj = ubi_file(path, block_size, start_offset, end_offset)
335+
ufile_obj = ubi_file(fileobj, block_size, start_offset, end_offset)
334336

335337
# Create UBI object
336338
ubi_obj = ubi(ufile_obj)
@@ -344,6 +346,7 @@ def main():
344346
make_files(ubi_obj, outpath)
345347

346348
ufile_obj.close()
349+
fileobj.close()
347350

348351

349352
if __name__=='__main__':

‎ubireader/ubi_io.py

+11-6
Original file line numberDiff line numberDiff line change
@@ -52,11 +52,15 @@ class ubi_file(object):
5252
def __init__(self, path, block_size, start_offset=0, end_offset=None):
5353
self.__name__ = 'UBI_File'
5454
self.is_valid = False
55-
try:
56-
log(self, 'Open Path: %s' % path)
57-
self._fhandle = open(path, 'rb')
58-
except Exception as e:
59-
error(self, 'Fatal', 'Open file: %s' % e)
55+
self._fileobj_passed = hasattr(path, "read")
56+
if self._fileobj_passed:
57+
self._fhandle = path
58+
else:
59+
try:
60+
log(self, 'Open Path: %s' % path)
61+
self._fhandle = open(path, 'rb')
62+
except Exception as e:
63+
error(self, 'Fatal', 'Open file: %s' % e)
6064

6165
self._fhandle.seek(0,2)
6266
file_size = self.tell()
@@ -111,7 +115,8 @@ def _get_block_size(self):
111115
block_size = property(_get_block_size)
112116

113117
def close(self):
114-
self._fhandle.close()
118+
if not self._fileobj_passed:
119+
self._fhandle.close()
115120

116121
def seek(self, offset):
117122
self._fhandle.seek(offset)

0 commit comments

Comments
 (0)