#!/usr/bin/python3

import argparse
import json
import sys

try:
    import yaml
except ImportError:
    print("[-] Please install yaml using the following command: 'pip install pyyaml'.", file=sys.stderr)
    sys.exit(-1)


def print_feeds_list(feeds_config):
    options = list()
    index = 0
    for provider, feeds in sorted(feeds_config['providers'].items()):

        for feed, feed_info in feeds.items():

            options.append(dict(provider=provider, feed=feed))
            print("[%d] %s - %s" % (index, provider, feed))
            index += 1
    return options


def sanitize_bot_id(id):
    id = id.replace('.', '')
    id = id.replace('_', '-')
    id = id.replace(' ', '-')
    id = id.lower()
    return id


def get_option_selected(options):
    print()

    while True:
        selection = input("Select a feed: ")
        try:
            selection = int(selection)
            if selection < len(options):
                break
        except Exception:
            print("Bad option.")
    return selection


def gen_configurations(provider, feed):
    runtime_config = dict()
    pipeline_config = dict()

    for bot_type, info in feeds_config['providers'][provider][feed]['bots'].items():

        bot_id = sanitize_bot_id(provider + "-" + feed + "-" + bot_type)

        if info['parameters']:
            parameters = info['parameters']
        else:
            parameters = dict()

        if bot_type == "collector":
            parameters["name"] = feed
            parameters["provider"] = provider

        runtime_config[bot_id] = dict(description="N/A",
                                      group=bot_type.title(),
                                      module=info['module'],
                                      parameters=parameters)

    collector_bot_id = sanitize_bot_id(provider + "-" + feed + "-collector")
    parser_bot_id = sanitize_bot_id(provider + "-" + feed + "-parser")

    pipeline_config[collector_bot_id] = dict()
    pipeline_config[collector_bot_id]['destination-queues'] = [parser_bot_id + "-queue"]

    pipeline_config[parser_bot_id] = dict()
    pipeline_config[parser_bot_id]['source-queue'] = parser_bot_id + "-queue"
    pipeline_config[parser_bot_id]['destination-queues'] = ["FIXME"]

    return runtime_config, pipeline_config


if __name__ == "__main__":

    parser = argparse.ArgumentParser(
        description='IntelMQ Feeds Config Generator tool',
    )

    parser.add_argument('--feeds-file',
                        action="store",
                        dest="feeds_file",
                        metavar="<filepath>",
                        required=True,
                        help='feeds.yaml config file')

    parser.add_argument('--all',
                        action="store_true",
                        dest="all",
                        required=False,
                        default=False,
                        help='iterate through all feeds')

    parser.add_argument('--runtime-output-file',
                        action="store",
                        dest="runtime_file",
                        metavar="<filepath>",
                        required=False,
                        help="write runtime configuration to the given file (e.g. '/tmp/runtime.conf') instead of stdout")

    parser.add_argument('--pipeline-output-file',
                        action="store",
                        dest="pipeline_file",
                        metavar="<filepath>",
                        required=False,
                        help="write pipeline configuration to the given file (e.g. '/tmp/pipeline.conf') instead of stdout")

    arguments = parser.parse_args()

    with open(arguments.feeds_file) as fp:
        feeds_config = yaml.load(fp.read())

    runtime_config = dict()
    pipeline_config = dict()

    if arguments.all:

        for provider, feeds in feeds_config['providers'].items():

            for feed in feeds.keys():
                tmp_runtime_config, tmp_pipeline_config = gen_configurations(provider, feed)
                runtime_config.update(tmp_runtime_config)
                pipeline_config.update(tmp_pipeline_config)

    else:

        options = print_feeds_list(feeds_config)
        selection = get_option_selected(options)

        provider = options[selection]['provider']
        feed = options[selection]['feed']

        runtime_config, pipeline_config = gen_configurations(provider, feed)

    if arguments.runtime_file:
        with open(arguments.runtime_file, "w") as fp:
            json.dump(runtime_config, fp, indent=4, sort_keys=True, separators=(',', ': '))
        print("Runtime configuration written to: %s" % arguments.runtime_file)
    else:
        print("\nRuntime configuration:\n")
        print(json.dumps(runtime_config, indent=4, sort_keys=True, separators=(',', ': ')))

    if arguments.pipeline_file:
        with open(arguments.pipeline_file, "w") as fp:
            json.dump(pipeline_config, fp, indent=4, sort_keys=True, separators=(',', ': '))
        print("Pipeline configuration written to: %s" % arguments.pipeline_file)
    else:
        print("\nPipeline configuration:\n")
        print(json.dumps(pipeline_config, indent=4, sort_keys=True, separators=(',', ': ')))
