#
# Licensed to Dagda under one or more contributor
# license agreements. See the NOTICE file distributed with
# this work for additional information regarding copyright
# ownership. Dagda licenses this file to you under
# the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#

import os
import shutil
import tempfile
from tarfile import TarFile


# Gets malware included in a docker image
def get_malware_included_in_docker_image(docker_driver, image_name=None, container_id=None):
    temp_dir = _extract_filesystem_bundle(docker_driver=docker_driver,
                                          image_name=image_name,
                                          container_id=container_id)
    raw_output = _clamav_run(docker_driver=docker_driver, temp_dir=temp_dir)
    output = _parse_raw_output(raw_output=raw_output)
    _clean_up(temporary_dir=temp_dir)
    return output


# -- Private methods

# Prepare filesystem bundle
def _extract_filesystem_bundle(docker_driver, container_id=None, image_name=None):
    temporary_dir = tempfile.mkdtemp()
    # Get and save filesystem bundle
    if container_id is not None:
        data = docker_driver.get_docker_client().export(container=container_id).data
        name = container_id
    else:
        data = docker_driver.get_docker_client().get_image(image=image_name).data
        name = image_name.replace('/', '_')
    with open(temporary_dir + "/" + name + ".tar", "wb") as file:
        file.write(data)
    # Untar filesystem bundle
    tarfile = TarFile(temporary_dir + "/" + name + ".tar")
    tarfile.extractall(temporary_dir)
    os.remove(temporary_dir + "/" + name + ".tar")
    # Return
    return temporary_dir


# Run ClamAV
def _clamav_run(docker_driver, temp_dir):
    # Docker pull for ensuring the geekduck/clamav image
    docker_driver.docker_pull('geekduck/clamav')
    # Start container
    container_id = docker_driver.create_container('geekduck/clamav',
                                                  "clamscan -ri  " + temp_dir,
                                                  [
                                                      temp_dir
                                                  ],
                                                  docker_driver.get_docker_client().create_host_config(
                                                      binds=[
                                                          temp_dir + ':' + temp_dir + ':ro'
                                                      ]))
    docker_driver.docker_start(container_id)
    # Wait for clamav
    output = docker_driver.docker_logs(container_id, True, False, True)
    # Stop container
    docker_driver.docker_stop(container_id)
    # Clean up
    docker_driver.docker_remove_container(container_id)
    # Return
    return output


# Parse raw output
def _parse_raw_output(raw_output):
    output = []
    for line in raw_output.splitlines():
        if line.endswith(' FOUND'):
            m = {}
            splitted_info = line.split(' ')
            m['file'] = splitted_info[0][:-1]
            m['malware'] = splitted_info[1]
            output.append(m)
    return output


# Clean the temporary directory
def _clean_up(temporary_dir):
    shutil.rmtree(temporary_dir)
