#!/usr/bin/env python3
# vim: set ts=4 sw=4 ts=4 et :

import argparse
import logging
import time

import greynoise
import collections

import matplotlib.pyplot as plt

logging.basicConfig(level=logging.INFO)
log = logging.getLogger()


def parse_args():
    parser = argparse.ArgumentParser(description="Default")
    parser.add_argument("--debug", help="debug", action="store_true")
    parser.add_argument("-q", "--query", help="GNQL Query", required=True)
    return parser.parse_args()


def main():
    args = parse_args()

    log.info("Running {}".format(__file__))
    if args.debug:
        log.setLevel(logging.DEBUG)
        log.debug("Debug mode enabled")

    # profiling
    s = time.perf_counter()

    gn_api_client = greynoise.GreyNoise()

    tag_cve = dict()
    metadata = gn_api_client.metadata()
    for tag in metadata.get("metadata", []):
        tag_cve[tag["name"]] = tag["cves"]

    cve_counts = collections.defaultdict(lambda: 0)
    scroll_token = None
    while True:
        results = gn_api_client.query(f"{args.query}", scroll=scroll_token)
        log.info(f"Got {results['count']} results")

        tag_count = collections.defaultdict(lambda: 0)
        for ip_data in results.get("data", []):
            for tag_name in ip_data.get("tags", []):
                tag_count[tag_name] += 1

        scroll_token = results.get("scroll", None)

        for tag_name, tag_count in tag_count.items():
            for cve in tag_cve[tag_name]:
                cve_counts[cve] += tag_count

        if results["complete"]:
            break

    print(cve_counts)
    plt.rcdefaults()
    fig, ax = plt.subplots()
    ax.barh(list(cve_counts.keys()), cve_counts.values())
    ax.set_label("CVE Counts")
    ax.set_title(f"{args.query}")
    plt.show()

    elapsed = time.perf_counter() - s
    log.info(f"{__file__} executed in {elapsed:0.5f} seconds.")


if __name__ == "__main__":
    main()