import fnmatch
import hashlib
import json
import logging
import math
import multiprocessing
import os
import shlex
import signal
import subprocess
import sys
import time
from multiprocessing import Process
from queue import Empty
from subprocess import check_output

from androguard.core.analysis.analysis import Analysis
from androguard.core.bytecodes import apk, dvm

import database
import settings

ENTROPY_THRESHOLD = 7.2
PRESENT = 1
LOADING = 2

byte_list = bytes([i for i in range(0, 256)])


def check_libraries(call):
    """

    Parameters
    ----------
    call

    Returns
    -------

    """
    lib_classes = {
        'ljava/lang/system;': [
            'load',
            'loadlibrary'
        ],
        'ljava/lang/runtime;': [
            'loadlibrary',
            'load',
        ],
        'lcom/getkeepsafe/relinker;': [
            'loadlibrary'
        ],
    }
    function = call.name.lower()
    class_ = call.class_name.lower()
    if class_ in lib_classes.keys():
        if function in lib_classes[class_]:
            return call
    return None


def check_dex_class_loader(method, call):
    function = call.name.lower()
    class_ = call.class_name.lower()
    classes = {
        'ldalvik/system/basedexclassloader;', 'ldalvik/system/dexclassloader;',
        'ldalvik/system/inmemorydexclassloader;', 'ldalvik/system/pathclassloader;'
    }
    if class_ in classes and function == '<init>':
        signature = method.get_method().get_triple()
        sha256_hash = hashlib.sha256(method.get_method().get_raw()).hexdigest()
        return signature, sha256_hash
    return None


def check_class_loader(method, call):
    """

    Parameters
    ----------
    method
    call

    Returns
    -------

    """
    loader_methods = {'<init>', 'loadclass'}
    function = call.name.lower()
    class_ = call.class_name.lower()
    if class_ == 'ljava/lang/classloader;':
        if function in loader_methods:
            signature = method.get_method().get_triple()
            sha256_hash = hashlib.sha256(method.get_method().get_raw()).hexdigest()
            return signature, sha256_hash
    return None


def extract(apk_path):
    """

    Parameters
    ----------
    apk_path

    Returns
    -------

    """
    tmpdir = os.path.dirname(os.path.realpath(apk_path))
    subprocess.check_call(shlex.split(f'unzip -o -d apk -qq {apk_path}'), cwd=tmpdir)
    return os.path.join(tmpdir, 'apk')


def shannon_entropy(data):
    """
    Taken from https://github.com/trufflesecurity/truffleHog/blob/dev/truffleHog/truffleHog.py
    """
    if not data:
        return 0
    entropy = 0
    for x in byte_list:
        p_x = float(data.count(x)) / len(data)
        if p_x > 0:
            entropy += - p_x * math.log(p_x, 2)
    return entropy


def file_info(filename):
    """

    Parameters
    ----------
    filename

    Returns
    -------

    """
    with open(filename, 'rb') as file:
        content = file.read()
    sha256 = hashlib.sha256(content).hexdigest()
    entropy = shannon_entropy(content)
    size = os.path.getsize(filename)
    return entropy, sha256, size


class Worker(Process):

    def __init__(self, name, in_queue, out_queue):
        super(Worker, self).__init__(name=f'[Analysis] Worker_{name}')
        self.name = name
        self.logger = settings.get_logger(f'worker_{name}')
        self.logger.setLevel(logging.DEBUG)
        self.input = in_queue
        self.results = out_queue
        self.run_lock = multiprocessing.RLock()
        self.run_status = True
        signal.signal(signal.SIGTERM, self.shutdown)
        signal.signal(signal.SIGINT, self.shutdown)

    def shutdown(self, signum, frame):
        with self.run_lock:
            self.run_status = False
            self.logger.info(f'Shutting down after current apk')

    def check_permissions(self, application, implied=False):
        """Checks the permissions the app uses.
    
        Lists only explicitly defined permissions unless implied is set.
    
        Parameters
        ----------
        application : apk.APK
            The application to check, as an androguard object.
        implied : bool
            If set, implied permissions will be checked as well.
        Returns
        -------
        list
            A deduplicated list of permissions used by the application.
        """
        permissions = list(set(application.get_permissions()))
        if implied:
            self.logger.debug('Including implied permissions')
            permissions += list(set(application.get_uses_implied_permission_list()))
        return permissions

    def check_packer(self, apk_path):
        """Uses APKiD to detect whether an apk was packed.
        Parameters
        ----------
        apk_path : str
            The application to run the detection on.
        Returns
        -------
        str or None:
            The name of the packer or None
        """
        try:
            output = check_output(['apkid', apk_path]).decode('UTF-8').split('\n')
        except Exception as e:
            error = repr(e)
            self.logger.error(f'ApkId encountered an error while processing {apk_path.split("/")[-1]}')
            return f'ERROR:\t{error}'
        for line in output:
            if 'packer :' in line:
                packer = line.split('packer :')[1].strip()
                self.logger.log(settings.VERBOSE, f'Application packed with {packer}')
                return packer
        self.logger.debug('No packer found')
        return None

    def check_reflection(self, call):
        """

        Parameters
        ----------
        call

        Returns
        -------

        """
        reflection_methods = ['getclass', 'class']
        reflection_classes = {
            'ljava/lang/classloader;': [
                '<init>',
                'loadclass'
            ],
            'ljava/lang/class;': [
                'forname'
            ]
        }
        function = call.name.lower()
        class_ = call.class_name.lower()
        if function in reflection_methods or class_.startswith('ljava/lang/reflect'):
            self.logger.debug(f'Found reflection using present class: {class_}:{function}')
            return PRESENT
        if class_ in reflection_classes.keys():
            if function in reflection_classes[class_]:
                self.logger.debug(f'Found reflection loading a new class: {class_}:{function}')
                return LOADING
        return False

    def check_files(self, application, apk_path):
        """Checks for the presence of native libraries.

        Parameters
        ----------
        application : apk.APK
            The application to check, as an androguard object.
        apk_path : str

        Returns
        -------
        dict
            A mapping from filename to sha256-hash for any native code
        dict
            A mapping from filename to entropy tuple entropy_value, hash
        """
        file_mapping = application.get_files_types()
        file_types = list(file_type.strip() for file_type in set(file_mapping.values()) if file_type.strip().
                          startswith('ELF'))
        file_type_string = '\t' + '\n\t'.join(file_types)
        extracted = None
        libraries = []
        encrypted = []
        for filename, file_type in file_mapping.items():
            if file_type.startswith('ELF'):
                name = filename.split(os.path.sep)[-1]
                if not extracted:
                    extracted = extract(apk_path)
                entropy, sha256, size = file_info(os.path.join(extracted, filename))
                self.logger.debug(f'Entropy: {entropy}\tSha: {sha256}')
                libraries.append((sha256, name, entropy, size))
            elif (file_type == 'Unknown' or file_type.startswith('Binary') or file_type == 'data') \
                    and filename != 'META-INF/CERT.RSA':
                if not extracted:
                    extracted = extract(apk_path)
                entropy, sha256, size = file_info(os.path.join(extracted, filename))
                if entropy > ENTROPY_THRESHOLD:
                    self.logger.log(settings.VERBOSE,
                                    f'Found potentially encrypted file: {file_type}\t{filename}\t{entropy}')
                    encrypted.append((sha256, filename, entropy, size))
        self.logger.debug(f'Found the following file types:\n{file_type_string}\n')
        return libraries, encrypted

    def check_methods(self, application):
        """

        Parameters
        ----------
        application

        Returns
        -------

        """
        analysis = Analysis()
        for dex in application.get_all_dex():
            analysis.add(dvm.DalvikVMFormat(dex))
        analysis.create_xref()
        reflection_access = 0
        reflection_code_loading = False
        class_loading = False
        dex_class_loading = False
        for class_name in analysis.classes:
            class_ = analysis.classes[class_name]
            for method in class_.get_methods():
                self.logger.debug(f'\tIn {method.name}')
                for _, call, _ in method.get_xref_to():
                    reflection = self.check_reflection(call)
                    if reflection:
                        if reflection == PRESENT:
                            reflection_access += 1
                        elif reflection == LOADING:
                            reflection_code_loading = True
                    elif check_class_loader(method, call):
                        class_loading = True
                    elif check_dex_class_loader(method, call):
                        dex_class_loading = True
#                    elif check_libraries(call):
#                        libraries.add(call)
        return reflection_access, reflection_code_loading, class_loading, dex_class_loading

    def run(self):
        run = True
        sleep = 1
        while run:
            try:
                inp = self.input.get_nowait()
                self.analyze(inp)
                sleep = 1
            except Empty:
                time.sleep(sleep)
                sleep = sleep * 2
            with self.run_lock:
                run = self.run_status
        sys.exit(0)

    def analyze(self, entry):
        """

        Parameters
        ----------
        entry

        Returns
        -------

        """
        tmpdir = None
        try:
            sha256, tmpdir = entry
            tmpdir = os.path.abspath(tmpdir)
            apks = fnmatch.filter(os.listdir(tmpdir), '*.apk')
            if len(apks) < 1:
                dir_str = '\n\t'.join([str(file) for file in os.listdir(tmpdir)])
                self.logger.error(f'No apks found for {sha256} in {tmpdir}\n\tFound instead:\n\t{dir_str}')
            apk_path = os.path.join(tmpdir, apks[0])
            self.logger.info(f'Analyzing {sha256}')
            application = apk.APK(apk_path)
            libraries, encrypted = self.check_files(application, apk_path)
            database.store_result(sha256, self.check_permissions(application), self.check_packer(apk_path),
                                  self.check_methods(application), libraries, encrypted)
        except Exception as e:
            error = repr(e)
            self.logger.error(f'Failed to analyze {sha256}.\n'
                              f'\tCause:\n{error}')
            database.access("INSERT INTO results (sha256, vt) VALUES (%s, %s)", (sha256, json.dumps({'error': error})))
            entry = (None, tmpdir)
        self.results.put(entry)
