import importlib
import os
import glob
import sys
import types
import malcat
from anomalies.base import Anomaly


SKIP_ANOMALY_IF_MORE_THAN_ELEMENTS = 200000
ANOMALY_MAX_LOCATIONS = 256
DEBUG_ANOMALY_TIMES = False

def get_scanners(data_dir, user_data_dir, reload_plugins=False):
    scanners = []
    seen = set()
    if user_data_dir and user_data_dir != data_dir:
        toscan = (user_data_dir, data_dir)
    else:
        toscan = (data_dir,)
    for plugins_path in toscan:
        for fname in glob.glob(os.path.join(plugins_path, "anomalies", "*.py")):
            module_name = os.path.relpath(fname, plugins_path)[:-3].replace(os.sep, ".")
            if module_name in sys.modules:
                mod = sys.modules[module_name]
                if reload_plugins:
                    mod = importlib.reload(mod)
            else:
                mod = importlib.import_module(module_name)
            for elemname in dir(mod):
                elem = getattr(mod, elemname)
                try:
                    if issubclass(elem, Anomaly) and elem != Anomaly and not elem in scanners:
                        if not elemname in seen:
                            elem.path = os.path.abspath(fname)
                            scanners.append(elem)
                            seen.add(elemname)
                except TypeError: pass
    return scanners


class AnomalyBinding:

    def __init__(self, scanner):
        self.scanner = scanner
        self.anomaly = malcat.Anomaly(scanner.__class__.__name__, scanner.__class__.level, scanner.__class__.category, scanner.__class__.__doc__, scanner.__class__.path)
        self.empty = True

    def add_location(self, address, size):
        if address is not None and size:
            self.anomaly.add_location(address, size)
        self.empty = False


def scan(scanners, analysis):
    todo = []
    if DEBUG_ANOMALY_TIMES:
        import datetime
        times = {}
    for scanner_class in scanners:
        if scanner_class.filetype:
            if type(scanner_class.filetype) in (list, tuple):
                if analysis.type not in scanner_class.filetype:
                    continue
            else:
                 if scanner_class.filetype != analysis.type:
                     continue
        if scanner_class.architecture:
            if type(scanner_class.architecture) in (list, tuple):
                if analysis.architecture not in scanner_class.architecture:
                    continue
            else:
                 if scanner_class.architecture != analysis.architecture:
                     continue
        scanner = scanner_class()
        anom = AnomalyBinding(scanner)
        todo.append(anom)
        if hasattr(scanner, "scan"):
            if DEBUG_ANOMALY_TIMES:
                start = datetime.datetime.now()
            try:
                iterator = iter(scanner.scan(analysis))
            except TypeError:
                raise ValueError("Wrong return type for Anomaly scanner {}: {}".format(scanner_class.__name__, type(scanner.scan(analysis))))
            for i, location in enumerate(iterator):
                if i >= ANOMALY_MAX_LOCATIONS:
                    break
                anom.add_location(*location)
            if DEBUG_ANOMALY_TIMES:
                times[anom.scanner.__class__.__name__] = datetime.datetime.now() - start
    # iterator optimisations
    for annotation_name, scan_function in [
            ("struct", "scan_structure"),
            ("fns", "scan_function"),
            ("strings", "scan_string"),
            ("carved", "scan_object"),
            ("cfg", "scan_bb"),
            ]:
        if not hasattr(analysis, annotation_name):
            continue
        todo_want_iterator = [(anom, getattr(anom.scanner, scan_function)) for anom in todo if hasattr(anom.scanner, scan_function)]
        if not todo_want_iterator:
            continue
        annotation = getattr(analysis, annotation_name)
        if len(annotation) > SKIP_ANOMALY_IF_MORE_THAN_ELEMENTS:
            print("Skipping {} anomaly evaluation because more than {} {} ({})".format(annotation_name, SKIP_ANOMALY_IF_MORE_THAN_ELEMENTS, annotation_name, len(annotation)))
            continue
        for element in annotation:
            for anom, function in todo_want_iterator:
                if DEBUG_ANOMALY_TIMES:
                    start = datetime.datetime.now()
                for i, location in enumerate(function(analysis, element)):
                    anom.add_location(*location)
                    if i >= ANOMALY_MAX_LOCATIONS:
                        break
                if DEBUG_ANOMALY_TIMES:
                    times[anom.scanner.__class__.__name__] = times.get(anom.scanner.__class__.__name__, datetime.timedelta()) + (datetime.datetime.now() - start)
    if DEBUG_ANOMALY_TIMES:
        s = sorted(times.items(), key = lambda x: -x[1])
        for i in s[:16]:
            print(i)
    # add all non-empty annotation
    return [a.anomaly for a in todo if not a.empty]
