#!/usr/bin/python

import collections
import config
import datetime
import db
import json
import misc
import models
import threading

from sqlalchemy import distinct

Counter = collections.namedtuple("Counter", ["cves_count", "urls_count"])

class SerializableBase:
    attr_names = []
    def __init__(self, **kwargs):
        for attr_name in self.attr_names:
            setattr(self, attr_name, kwargs.get(attr_name))

    def serialize_to_json(self):
        return {name: getattr(self, name) for name in self.attr_names}

class FuncLaunch(SerializableBase):
    attr_names = [
        "date_of_launch", # as string, not as datetime.datetime!
        "elapsed_time", # as string, not as datetime.timedelta!
        "new_unique_cves",
        "new_url_matches",
        "func_name",
    ]

class Stats(SerializableBase):
    attr_names = [
        "funcs",
        "comment",
        "total_unique_cves",
        "total_url_matches",
        "id",
    ]

    def add_func_launch(self, func):
        if not self.funcs:
            self.funcs = []
            self.total_unique_cves = 0
            self.total_url_matches = 0
        self.funcs.append(func)
        self.total_unique_cves += func.new_unique_cves
        self.total_url_matches += func.new_url_matches

def measure_stats():
    count = lambda x: db.global_session.query(distinct(x)).count()

    cve_count = \
        count(list(models.cve_url_table.columns)[0]) + \
        count(list(models.non_cve_url_table.columns)[0])

    url_count = \
        count(list(models.cve_url_table.columns)[1]) + \
        count(list(models.non_cve_url_table.columns)[1])

    counter = Counter(cve_count, url_count)
    return counter


def record_stats_decorator(func):
    def wrapper(*args, **kwargs):
        # Measure the execution stats.
        before = measure_stats()

        print "{0} started.".format(func.func_name)
        start = datetime.datetime.now()
        func(*args, **kwargs)
        delta = datetime.datetime.now() - start
        print "{0} ended.".format(func.func_name)
        print "=" * 80

        after = measure_stats()

        # Read previous data.
        try:
            with open(config.STATS_FILE, "r") as f:
                prev_stats = json.load(f)
        except:
            # If the file does not exists or is empty.
            prev_stats = []

        # Marshal the data.
        first_func_launch = not hasattr(config, config.STATS_DECORATOR_KEY)
        if first_func_launch:
            current_stats = Stats()
            setattr(config, config.STATS_DECORATOR_KEY, current_stats)
            current_id = reduce(max, [x["id"] for x in prev_stats], -1) + 1
            current_stats.id = current_id
        else:
            current_stats = getattr(config, config.STATS_DECORATOR_KEY)
        current_func_launch = FuncLaunch(
                date_of_launch=start.strftime(config.DATE_TIME_FORMAT),
                elapsed_time=str(delta),
                new_unique_cves=after.cves_count - before.cves_count,
                new_url_matches=after.urls_count - before.urls_count,
                func_name=func.func_name
            )
        current_stats.add_func_launch(current_func_launch)

        # Combine current and previous data.
        if first_func_launch:
            prev_stats.append(current_stats)
        else:
            # Fields are sorted by id.
            prev_stats[current_stats.id] = current_stats

        # Write combined data.
        with open(config.STATS_FILE, "w") as f:
            json.dump(prev_stats, f, indent=4,
                    default=SerializableBase.serialize_to_json)

    return wrapper

def count_cves_by_vendor():
    vendor_dict = count_cves_by_vendor.vendor_dict = {}
    def item_handler(item):
        vendors = item["cve"]["affects"]["vendor"]["vendor_data"]
        for v in vendors:
            name = v["vendor_name"]
            count_cves_by_vendor.vendor_dict.setdefault(name, 0)
            count_cves_by_vendor.vendor_dict[name] += 1
    misc.crawl_nist_files(item_handler)
    vendor_list = sorted(vendor_dict.iteritems(), key=lambda x: x[1], reverse=True)
    answer = []
    cur_count = 0
    total_count = sum([x[1] for x in vendor_list])
    for v in vendor_list:
        cur_count += v[1]
        answer.append({
            "name": v[0],
            "count": v[1],
            "sum_till_here": cur_count,
            "ratio_till_here": "{0:.2f}%".format(cur_count * 100.0 / total_count),
        })

    print "total number of vendors: {0}".format(len(vendor_list))
    print "total number of cves: {0}".format(total_count)

    with open(config.VENDORS_FILE, "w") as f:
        json.dump(answer, f, indent=4)

