Added additional file access abstraction to common.BlockFile; fixed some Python3…

Added additional file access abstraction to common.BlockFile; fixed some Python3 compatability bugs.

src/binwalk/common.py

@@ -93,8 +93,26 @@ def unique_file_name(base_name, extension=''):
 
 class BlockFile(io.FileIO):
 	'''
-	Abstraction class to handle reading data from files in blocks.
-	Necessary for large files.
+	Abstraction class for accessing binary files.
+
+	This class overrides io.FilIO's read and write methods. This guaruntees two things:
+
+		1. All requested data will be read/written via the read and write methods.
+		2. All reads return a str object and all writes can accept either a str or a
+		   bytes object, regardless of the Python interpreter version.
+
+	However, the downside is that other io.FileIO methods won't work properly in Python 3,
+	namely things that are wrappers around self.read (e.g., readline, readlines, etc).
+
+	This class also provides a read_block method, which is used by binwalk to read in a
+	block of data, plus some additional data (MAX_TRAILING_SIZE), but on the next block read
+	pick up at the end of the previous data block (not the end of the additional data). This
+	is necessary for scans where a signature may span a block boundary.
+
+	The descision to force read to return a str object instead of a bytes object is questionable
+	for Python 3, it seemed the best way to abstract differences in Python 2/3 from the rest
+	of the code (especially for people writing plugins) and to add Python 3 support with 
+	minimal code change.
 	'''
 
 	# The MAX_TRAILING_SIZE limits the amount of data available to a signature.
@@ -150,12 +168,51 @@ class BlockFile(io.FileIO):
 
 		self.seek(self.offset)
 			
+	def write(self, data):
+		'''
+		Writes data to the opened file.
+		
+		io.FileIO.write does not guaruntee that all data will be written;
+		this method overrides io.FileIO.write and does guaruntee that all data will be written.
+
+		Returns the number of bytes written.
+		'''
+		n = 0
+		l = len(data)
+		data = str2bytes(data)
+
+		while n < l:
+			n += io.FileIO.write(self, data[n:])
+
+		return n
+
+	def read(self, n=-1):
+		''''
+		Reads up to n bytes of data (or to EOF if n is not specified).
+		
+		io.FileIO.read does not guaruntee that all requested data will be read;
+		this method overrides io.FileIO.read and does guaruntee that all data will be read.
+
+		Returns a str object containing the read data.
+		'''
+		l = 0
+		data = b''
+
+		while n < 0 or l < n:
+			tmp = io.FileIO.read(self, n-l)
+			if tmp:
+				data += tmp
+				l += len(tmp)
+			else:
+				break
+
+		return bytes2str(data)
 
 	def read_block(self):
 		'''
 		Reads in a block of data from the target file.
 
-                Returns a tuple of (file block data, block data length).
+                Returns a tuple of (str(file block data), block data length).
                 '''
 		dlen = 0
 		data = None
@@ -167,7 +224,6 @@ class BlockFile(io.FileIO):
 			data = self.read(self.READ_BLOCK_SIZE + self.MAX_TRAILING_SIZE)
 
 			if data and data is not None:
-				data = bytes2str(data)
 
 				# Get the actual length of the read in data
 				dlen = len(data)

src/binwalk/extractor.py

@@ -176,8 +176,9 @@ class Extractor:
 		'''
 		try:
 			# Process each line from the extract file, ignoring comments
-			for rule in open(fname).readlines():
-				self.add_rule(rule.split(self.COMMENT_DELIM, 1)[0])
+			with open(fname, 'r') as f:
+				for rule in f.readlines():
+					self.add_rule(rule.split(self.COMMENT_DELIM, 1)[0])
 		except Exception as e:
 			raise Exception("Extractor.load_from_file failed to load file '%s': %s" % (fname, str(e)))
 

src/binwalk/parser.py

@@ -140,8 +140,7 @@ class MagicParser:
 		line_count = 0
 
 		try:
-			for line in io.FileIO(file_name).readlines():
-				line = bytes2str(line)
+			for line in open(file_name, 'r').readlines():
 				line_count += 1
 
 				# Check if this is the first line of a signature entry

src/binwalk/plugins.py

@@ -126,7 +126,7 @@ class Plugins:
 				if val is not None:
 					retval |= val
 			except Exception as e:
-				sys.stderr.write("WARNING: %s.%s failed: %s\n" % (str(callback.im_class), callback.__name__, str(e)))
+				sys.stderr.write("WARNING: %s.%s failed: %s\n" % (callback.__module__, callback.__name__, e))
 
 		return retval
 

src/binwalk/plugins/lzmamod.py

@@ -31,7 +31,7 @@ class Plugin:
 		# Try extracting the LZMA file without modification first
 		if not self.binwalk.extractor.execute(self.original_cmd, fname):
 			out_name = os.path.splitext(fname)[0] + '-patched' + os.path.splitext(fname)[1]
-			fp_out = open(out_name, 'wb')
+			fp_out = BlockFile(out_name, 'w')
 			fp_in = BlockFile(fname)
 			fp_in.MAX_TRAILING_SIZE = 0
 			i = 0
@@ -44,7 +44,7 @@ class Plugin:
 				else:
 					out_data = data
 				
-				fp_out.write(str2bytes(out_data))
+				fp_out.write(out_data)
 	
 				i += dlen
 

src/binwalk/prettyprint.py

@@ -4,7 +4,6 @@ import hashlib
 import csv as pycsv
 from datetime import datetime
 from binwalk.compat import *
-from binwalk.common import BlockFile
 
 class PrettyPrint:
 	'''
@@ -134,7 +133,7 @@ class PrettyPrint:
 		'''
 		md5 = hashlib.md5()
                 
-		with BlockFile(file_name, 'r') as f:
+		with open(file_name, 'rb') as f:
 			for chunk in iter(lambda: f.read(128*md5.block_size), b''):
 				md5.update(chunk)