#!/usr/bin/env python

import os
import sys
import glob
import json
import pkgutil
import logging
import argparse
import importlib
from mwdblib import MWDB
from functools import partial
from multiprocessing import Pool
from configparser import ConfigParser

log = logging.getLogger(__name__)

__author__  = "c3rb3ru5"
__version__ = "1.0.0"

def mwdb_feeds_module_worker(config, mwdb, module):
    spec = importlib.util.spec_from_file_location("module.name", module)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    module = module.MWDBFeedsModule(config=config, mwdb=mwdb)
    if module.enabled is True:
        return module.main()

class MWDBFeeds():

    """
    A Modular MWDB Collector
    """

    def __init__(self):
        self.name = 'mwdb-feeds'

    def get_modules(self):
        modules = glob.glob('{modules}/**/*.py'.format(
                modules=self.config.get(self.name, 'modules')
            ),
            recursive=True
        )
        return modules

    def arguments(self):
        self.parser = argparse.ArgumentParser(
            prog=f'mwdb-feeds v{__version__}',
            description='A Module Malware Collector for MWDB',
            epilog=f'Author: {__author__}'
        )
        self.parser.add_argument(
            '--version',
            action='version',
            version=f'v{__version__}'
        )
        self.parser.add_argument(
            '-c',
            '--config',
            default=None,
            help='Configuration File',
            required=True
        )
        self.parser.add_argument(
            '-m',
            '--module',
            default=False,
            help='Execute Single Module',
            required=False
        )
        self.parser.add_argument(
            '-l',
            '--list-modules',
            default=False,
            action='store_true',
            required=False
        )
        self.args = self.parser.parse_args()
    
    def check_config(self):
        options = ['threads', 'log', 'modules', 'debug', 'url', 'username', 'api_key', 'verify_ssl']
        for option in options:
            if self.config.has_option(self.name, option) is False:
                print(f'{self.name} section is missing the {option} option')
                self.parser.print_help()
                sys.exit(1)

    def setup_logging(self):
        level = logging.INFO
        if self.config.getboolean(self.name, 'debug'):
            level = logging.DEBUG
        if self.config.get(self.name, 'log') not in ['False', False, None, '']:
            logging.basicConfig(filename=self.config.get(self.name, 'log'), level=level)
        else:
            logging.basicConfig(level=level)

    @staticmethod
    def list_modules(modules):
            for module in modules:
                print(os.path.basename(module)[:-3])
            sys.exit(0)

    @staticmethod
    def get_module(single_module, modules):
        for module in modules:
            if os.path.basename(module)[:-3] == single_module:
                return module
        return False

    def run_single_module(self, modules):
        if self.args.module:
            module = self.get_module(self.args.module, modules)
            if module:
                result = mwdb_feeds_module_worker(self.config, self.mwdb, module)
                if result is not None:
                    result = [result]
                    print(json.dumps(result, indent=4))
                self.logout()
                sys.exit(0)
            else:
                print(f'the module {self.args.module} does not exist', file=sys.stderr)
                self.parser.print_help()
                sys.exit(1)

    def run_list_modules(self, modules):
        if self.args.list_modules:
            self.list_modules(modules)
            sys.exit(0)

    def run_all_modules(self, modules):
        pool = Pool(processes=self.config.getint(self.name, 'threads'))
        results = pool.map(partial(mwdb_feeds_module_worker, self.config, self.mwdb), modules)
        results = list(filter(None, results))
        print(json.dumps(results, indent=4))
        sys.exit(0)

    def login(self) -> None:
        try:
            self.mwdb = MWDB()
            self.mwdb.api.api_url = self.config.get(self.name, 'url')
            self.mwdb.api.verify_ssl = self.config.getboolean(self.name, 'verify_ssl')
            self.mwdb.api.set_api_key(self.config.get(self.name, 'api_key'))
        except Exception as error:
            log.error(error)
            sys.exit(1)

    def logout(self):
        self.mwdb.logout()

    def main(self):
        self.arguments()
        self.config = ConfigParser()
        self.config.read(self.args.config)
        self.check_config()
        self.setup_logging()
        self.login()
        modules = self.get_modules()
        self.run_list_modules(modules)
        self.run_single_module(modules)
        self.run_all_modules(modules)
        self.logout()

if __name__ in '__main__':
    try:
        collector = MWDBFeeds()
        collector.main()
    except KeyboardInterrupt:
        log.warning('interrupted, exiting...')
        sys.exit(0)
    except Exception as error:
        log.error(error)
        sys.exit(1)
