import zlib
from binwalk.plugins import *

class Plugin:
	'''
	Finds and validates zlib compressed data.
	'''

	ENABLED = False

	def __init__(self, binwalk):
		self.fd = None
		self.enabled = False

		if binwalk.scan_type == binwalk.BINWALK:
			# Add the zlib file to the list of magic files
			binwalk.magic_files.append(binwalk.config.paths['system'][binwalk.config.ZLIB_MAGIC_FILE])
			self.enabled = True
	
	def pre_scan(self, fd):
		if self.enabled:
			self.fd = open(fd.name, 'rb')

	def callback(self, result):

		# If this result is a zlib signature match, try to decompress the data
		if self.fd and result['description'].lower().startswith('zlib'):
			try:
				# With a lot of zlib matches on large files, this can result in a lot of disk I/O
				# and a lot of memory usage. This is why this module is disabled by default.
				self.fd.seek(result['offset'])
				zlib.decompress(self.fd.read())
			except:
				# If decompression raised an exception, this is probably not valid zlib data
				return (PLUGIN_NO_DISPLAY | PLUGIN_NO_EXTRACT)
		
		return PLUGIN_CONTINUE

	def post_scan(self, fd):
		if self.fd:
			self.fd.close()

