# Copyright 2021 Google LLC
#
# Licensed 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.
''' Uses the SymCC-AFL hybrid from SymCC. '''

import os
import time
import shutil
import threading
import subprocess

from fuzzers import utils
from fuzzers.afl import fuzzer as afl_fuzzer
from fuzzers.aflplusplus import fuzzer as aflplusplus_fuzzer


def get_symcc_build_dir(target_directory):
    """Return path to uninstrumented target directory."""
    return os.path.join(target_directory, 'uninstrumented')


def build():
    """Build an AFL version and SymCC version of the benchmark"""
    print('Step 1: Building with AFL and SymCC')
    build_directory = os.environ['OUT']

    # First build with AFL.
    src = os.getenv('SRC')
    work = os.getenv('WORK')
    with utils.restore_directory(src), utils.restore_directory(work):
        # Restore SRC to its initial state so we can build again without any
        # trouble. For some OSS-Fuzz projects, build_benchmark cannot be run
        # twice in the same directory without this.
        aflplusplus_fuzzer.build('tracepc', 'symcc')

    print('Step 2: Completed AFL build')
    # Copy over AFL artifacts needed by SymCC.
    shutil.copy('/afl/afl-fuzz', build_directory)
    shutil.copy('/afl/afl-showmap', build_directory)

    # Copy over symcc artifacts and symbolic libc++.
    print('Step 3: Copying SymCC files')
    symcc_build_dir = get_symcc_build_dir(os.environ['OUT'])
    shutil.copy(
        '/symcc/build//SymRuntime-prefix/src/SymRuntime-build/libSymRuntime.so',
        symcc_build_dir)
    shutil.copy('/usr/lib/libz3.so', os.path.join(symcc_build_dir, 'libz3.so'))
    shutil.copy('/libcxx_native_build/lib/libc++.so.1', symcc_build_dir)
    shutil.copy('/libcxx_native_build/lib/libc++abi.so.1', symcc_build_dir)
    shutil.copy('/rust/bin/symcc_fuzzing_helper', symcc_build_dir)


def launch_afl_thread(input_corpus, output_corpus, target_binary,
                      additional_flags):
    """ Simple wrapper for running AFL. """
    afl_thread = threading.Thread(target=afl_fuzzer.run_afl_fuzz,
                                  args=(input_corpus, output_corpus,
                                        target_binary, additional_flags))
    afl_thread.start()
    return afl_thread


def fuzz(input_corpus, output_corpus, target_binary):
    """
    Launches a master and a secondary instance of AFL, as well as
    the symcc helper.
    """
    target_binary_dir = os.path.dirname(target_binary)
    symcc_workdir = get_symcc_build_dir(target_binary_dir)
    target_binary_name = os.path.basename(target_binary)
    symcc_target_binary = os.path.join(symcc_workdir, target_binary_name)

    os.environ['AFL_DISABLE_TRIM'] = '1'

    # Start a master and secondary instance of AFL.
    # We need both because of the way SymCC works.
    print('[run_fuzzer] Running AFL for SymCC')
    afl_fuzzer.prepare_fuzz_environment(input_corpus)
    launch_afl_thread(input_corpus, output_corpus, target_binary,
                      ['-S', 'afl-secondary'])
    time.sleep(5)

    # Start an instance of SymCC.
    # We need to ensure it uses the symbolic version of libc++.
    print('Starting the SymCC helper')
    new_environ = os.environ.copy()
    new_environ['LD_LIBRARY_PATH'] = symcc_workdir
    cmd = [
        os.path.join(symcc_workdir,
                     'symcc_fuzzing_helper'), '-o', output_corpus, '-a',
        'afl-secondary', '-n', 'symcc', '-m', '--', symcc_target_binary, '@@'
    ]
    with subprocess.Popen(cmd, env=new_environ):
        pass
