#!/usr/bin/python

import config
import db
import fs
import json
import misc
import models
import net
import os
import re
import stats

from os.path import basename, splitext
from threading import Thread, Event
from Queue import Queue


urls_re_whitelist = [
        re.compile(r"https?://github.com\S+"),
        re.compile(r"https?://svn.apache.org\S+"),
        re.compile(r"https?://git.kernel.org\S+"),
        re.compile(r"https?://\S+.googlesource.com\S+"),
        re.compile(r"https?://codereview.chromium.org\S+"),
        re.compile(r"https?://.+?commit\S+")
]

def _should_keep_url(url):
    for url_re in urls_re_whitelist:
        if re.match(url_re, url):
            return True
    return False

""" Returns a dict {cve_string:cve} for all cves present in the DB. """
def _index_cves():
    cves = db.global_session.query(models.CVE).all()
    return {cve.cve_string:cve for cve in cves}

"""
    Returns a list of all queries necessary to add all urls in url_strings,
connect them to the cve corresponding to cve_string, and optionally, add
tags to the cve.
"""
def _get_queries_for_cve_url(cve_string, url_strings, tags=[]):
    queries = \
        [db.InsertQuery(models.URL, url=url_string)
            for url_string in url_strings
        ] + \
        [db.ConnectQuery(models.cve_url_table, cve_string, url_string)
            for url_string in url_strings
        ]
    if url_strings and tags:
        queries += [db.UpdateTagQuery(models.CVE, cve_string, tags)]
    return queries

""" Does the same but for nonCVE objects. """
def _get_queries_for_noncve_url(hash_id, url_strings, tags=[]):
    queries = \
        [db.InsertQuery(models.URL, url=url_string)
            for url_string in url_strings
        ] + \
        [db.ConnectQuery(models.non_cve_url_table, hash_id, url_string)
            for url_string in url_strings
        ]
    if url_strings and tags:
        queries += [db.UpdateTagQuery(models.NonCVE, hash_id, tags)]
    return queries

"""
    For each element in sequence invokes worker_func, gathers the resulting
queries, appends them in a list and passes it to db.process_queries. Optionally
finishes the bar object.
"""
def _process_queries_from_workers(worker_func, sequence, parallel, bar=None):
    results = misc.map_func(worker_func, sequence, parallel)
    if bar:
        bar.finish()
    queries = misc.unique(misc.flatten_list(results))
    return db.process_queries(queries)

@stats.record_stats_decorator
def crawl_nist_feed(parallel=False):
    # Works faster without parallelism.
    def items_handler(cve_items):
        get_cve = lambda item: item["cve"]["CVE_data_meta"]["ID"]
        get_urls = lambda item: item["cve"]["references"]["reference_data"]

        def get_cwe_strings(item):
            return [desc["value"]
                    for data in item["cve"]["problemtype"]["problemtype_data"]
                    for desc in data["description"]
                    ]

        queries = []
        for item in cve_items:
            queries.append(db.InsertQuery(models.CVE, cve_string=get_cve(item)))

            for url_string in filter(_should_keep_url,
                    [url["url"] for url in get_urls(item)]):
                queries.append(db.InsertQuery(models.URL, url=url_string))
                queries.append(db.ConnectQuery(models.cve_url_table,
                    get_cve(item), url_string))

            for cwe_string in get_cwe_strings(item):
                queries.append(db.InsertQuery(models.CWE, cwe_string=cwe_string))
                queries.append(db.ConnectQuery(models.cve_cwe_table,
                    get_cve(item), cwe_string))

        return queries

    result = misc.crawl_nist_files(items_handler, parallel=parallel)
    queries = misc.flatten_list(result)

    print "Done parsing files. Storing results into the DB"
    #TODO: populate cwes with descriptions

    db.process_queries(queries)
    misc.update_cwe_descriptions()

@stats.record_stats_decorator
def crawl_android_cve_checker(local):
    # Works faster with parallelism.
    api_url_t = "https://api.github.com/repos/torvalds/linux/commits/{0}"
    http_url_t = "https://github.com/torvalds/linux/commit/{0}"
    hash_re = re.compile(r"([a-f0-9]{40})")

    cve_index = _index_cves()
    def file_func(file_obj):
        name, ext = splitext(basename(file_obj["path"]))
        queries = []
        if ext == ".patch" and cve_index.get(name):
            if local:
                raw = fs.read_file(file_obj["path"])
            else:
                raw = net.get_raw_resource(file_obj["download_url"],
                    auth=net.github_auth)
            hashes = re.findall(hash_re, raw)
            for h in hashes:
                api_url = api_url_t.format(h)
                http_url = http_url_t.format(h)
                try:
                    net.get_json_resource(api_url, auth=net.github_auth)
                except:
                    # not upstream commit. Shit happens
                    # print "bad response for resource: " + api_url
                    pass
                else:
                    queries = _get_queries_for_cve_url(name, [http_url],
                            ["C", "Linux", "github"])
        bar.update()
        return ([], queries)

    def dir_func(dir_obj):
        if  "patches" in dir_obj["path"]:
            return True
        return False

    bar = misc.UnknownLengthBar(parallel=True)
    queries = crawl_github_repo(user="raymanfx", repo="android-cve-checker",
            file_callback=file_func, dir_callback=dir_func, local=local)
    bar.finish()
    return db.process_queries(queries)

@stats.record_stats_decorator
def crawl_linux_kernel_cves():
    # Works fast enough without parallelism.
    url = "https://raw.githubusercontent.com/nluedtke/linux_kernel_cves/master/stream_fixes.json"
    http_url_t = "https://github.com/torvalds/linux/commit/{0}"
    patches = net.get_json_resource(url)
    bar = misc.KnownLengthBar(maxval=len(patches), parallel=False)

    cve_index = _index_cves()

    def worker(item):
        cve_string = item[0]
        cve = cve_index.get(cve_string)
        queries = []
        if cve:
            url_strings = [http_url_t.format(v["cmt_id"]) for v in item[1].values()]
            queries = _get_queries_for_cve_url(cve_string, url_strings,
                    ["C", "Linux", "github"])
        bar.update()
        return queries

    return _process_queries_from_workers(worker, patches.items(), False, bar)


@stats.record_stats_decorator
def crawl_vuln_db(local, parallel=True):
    commit_re = re.compile("(https://github.com/\S+/commits?/[a-f0-9]+)")
    cve_re = re.compile("(CVE-\d+-\d+)")
    tag_dict = {
            "composer": "PHP",
            "golang"  : "Go",
            "maven"   : "Java",
            "npm"     : "JS",
            "nuget"   : "nuget", #has different languages
            "pip"     : "Python",
            "rubygems": "Ruby",
    }
    def file_func(file_obj):
        name = basename(file_obj["path"])
        if local:
            raw = fs.read_file(file_obj["path"])
        else:
            raw = net.get_raw_resource(file_obj["download_url"], auth=net.github_auth)
        url_strings = misc.unique(re.findall(commit_re, raw))
        queries = []
        if url_strings:
            # Figure out the tag list.
            tags = ["vulndb"]
            for k, v in tag_dict.iteritems():
                if k in file_obj["path"]:
                    tags.append(v)
                    break

            # Insert CVEs/NonCVEs and connect them to urls.
            cve_strings = misc.unique(re.findall(cve_re, raw))
            if not cve_strings:
                hash_id = models.NonCVE.hash_id_for_urls(url_strings)
                queries = [db.InsertQuery(models.NonCVE, hash_id=hash_id)]
                queries += _get_queries_for_noncve_url(hash_id, url_strings, tags)
            else:
                # Surpisingly there's some CVEs in this db, which are marked
                # as reserved in the nist feed. We need to add them here.
                queries = [db.InsertQuery(models.CVE, cve_string=cve_string)
                        for cve_string in cve_strings]
                for cve_string in cve_strings:
                    queries += _get_queries_for_cve_url(cve_string, url_strings, tags)
        bar.update()
        return ([], queries)

    def dir_func(dir_obj):
        if "data" in dir_obj["path"]:
            return True
        return False

    bar = misc.UnknownLengthBar(parallel=parallel)
    queries = crawl_github_repo(user="snyk", repo="vulnerabilitydb",
            file_callback=file_func, dir_callback=dir_func, local=local)
    bar.finish()

    return db.process_queries(queries)

@stats.record_stats_decorator
def crawl_django():
    cve_re = re.compile(r":cve:`(\S+)`")
    commit_re = re.compile(r"(https://github.com/\S+/[a-f0-9]+)")
    cve_index = _index_cves()
    raw = net.get_raw_resource(
        "https://raw.githubusercontent.com/django/django/master/docs/releases/security.txt")

    indices = [x.start() for x in re.finditer(r":cve:", raw)]
    sub_strings = [raw[s:f] for s,f in zip(indices, indices[1:] + [len(raw)])]

    bar = misc.KnownLengthBar(maxval=len(indices), parallel=False)
    def worker(sub_string):
        queries = []
        cve_string = "CVE-" + re.findall(cve_re, sub_string)[0]
        cve = cve_index.get(cve_string)
        if not cve:
            print "CVE not found?!: " + cve_string
            return []

        # Find the URLs
        url_strings = re.findall(commit_re, sub_string)
        if url_strings:
            queries = _get_queries_for_cve_url(cve_string, url_strings,
                    ["Python", "Django"])
        return queries
    return _process_queries_from_workers(worker, sub_strings, False, bar)

@stats.record_stats_decorator
def crawl_debian_security(local, parallel=True):
    # Works faster with parallelism.
    cve_re = re.compile("^(CVE\S+)")
    list_url = "https://salsa.debian.org/security-tracker-team/security-tracker/raw/master/data/CVE/list"

    if not local:
        raw = net.get_raw_resource(list_url)
    else:
        path = misc.repo_path("security-tracker-team", "security-tracker")
        path = os.path.join(path, "raw", "master", "data", "CVE", "list")
        with open(path, "r") as f:
            raw = f.read()

    indices = [x.start() for x in re.finditer(r"^CVE", raw, re.MULTILINE)]
    sub_strings = [raw[s:f] for s,f in zip(indices, indices[1:] + [len(raw)])]
    cve_index = _index_cves()

    bar = misc.KnownLengthBar(maxval=len(sub_strings), parallel=parallel)
    def worker(sub_string):
        cve_string = re.findall(cve_re, sub_string)[0]
        cve = cve_index.get(cve_string)
        queries = []
        if cve:
            url_strings = misc.flatten_list([re.findall(url_re, sub_string)
                    for url_re in urls_re_whitelist])
            queries = _get_queries_for_cve_url(cve_string, url_strings,
                    ["DebianSec"])
        bar.update()
        return queries

    return _process_queries_from_workers(worker, sub_strings, parallel, bar)

@stats.record_stats_decorator
def crawl_debian_fake_names(parallel=True):
    # Works faster with parallelism.
    main_page = net.get_raw_resource("https://security-tracker.debian.org/tracker/data/fake-names")
    temp_re = re.compile(r"/tracker/(TEMP-[0-9A-F-]+)")
    href_re = re.compile(r'href="(\S+?)"')

    temps = re.findall(temp_re, main_page)
    bar = misc.KnownLengthBar(maxval=len(temps), parallel=parallel)
    def worker(temp):
        raw = net.get_raw_resource("https://security-tracker.debian.org/tracker/{0}".format(temp))
        url_strings = [x for x in re.findall(href_re, raw) if _should_keep_url(x)]
        queries = [db.InsertQuery(models.NonCVE, hash_id=temp)]
        queries += _get_queries_for_noncve_url(temp, url_strings, ["DebianFake"])
        bar.update()
        return queries

    return _process_queries_from_workers(worker, temps, parallel, bar)

@stats.record_stats_decorator
def crawl_android_security_bulletin(parallel=True):
    raw = net.get_raw_resource("https://source.android.com/security/bulletin/")
    urls = re.findall(r"https?://source.android.com/security/bulletin/[0-9-]{10}", raw)
    urls = misc.unique(urls)
    bar = misc.KnownLengthBar(maxval=len(urls), parallel=parallel)

    def worker(url_string):
        raw = net.get_raw_resource(url_string)
        results = re.findall(r"<td>(CVE[0-9-]+)</td>\s+<td><a href=\"(\S+?)\">", raw, re.DOTALL)
        results = [r for r in results if _should_keep_url(r[1])]
        queries =  [db.InsertQuery(models.CVE, cve_string=r[0]) for r in results]
        queries += [db.InsertQuery(models.URL, url=r[1]) for r in results]
        queries += [db.ConnectQuery(models.cve_url_table, r[0], r[1]) for r in results]
        bar.update()
        return queries

    return _process_queries_from_workers(worker, urls, parallel, bar)

@stats.record_stats_decorator
def crawl_jenkins():
    url_t = "https://github.com/jenkinsci/jenkins/commit/{0}"
    cve_t = "JENKINS-{0}"
    commits = misc.get_repo_log("jenkinsci", "jenkins")
    hashes = [c["hash"] for c in commits if "SECURITY" in c["message"]]
    url_strings = [url_t.format(h) for h in hashes]
    queries = []
    for i, url_string in enumerate(url_strings):
        hash_id=cve_t.format(i)
        queries += [
            db.InsertQuery(models.NonCVE, hash_id=hash_id),
            db.InsertQuery(models.URL, url=url_string),
            db.ConnectQuery(models.non_cve_url_table, hash_id, url_string),
            db.UpdateTagQuery(models.NonCVE, hash_id, ["Jenkins", "Java"])
        ]
    db.process_queries(queries)

def crawl_github_repo(user, repo, file_callback, dir_callback, local):
    # TODO: add a non-parallel version.
    # TODO: adapt to the new db format. Maybe make a second queue for collecting
    # return data from workers.
    task_queue = Queue()
    results_queue = Queue()
    if local:
        path = misc.repo_path(user, repo)
        crawl_list = fs.crawl_dir(path)
    else:
        base_url = net.get_api_url(user, repo)
        crawl_list = net.get_json_resource(base_url, auth=net.github_auth)

    def worker(task_queue, results_queue):
        while True:
            item = task_queue.get()
            crawl_list = []
            if item["type"] == "file":
                crawl_list, queries = file_callback(item)
                results_queue.put(queries)
            elif dir_callback(item):
                if local:
                    crawl_list = fs.crawl_dir(item["path"])
                else:
                    crawl_list = net.get_json_resource(item["url"], auth=net.github_auth)
            for item in crawl_list: task_queue.put(item)
            task_queue.task_done()

    for item in crawl_list: task_queue.put(item)
    for i in xrange(0, config.THREADS_COUNT):
        t = Thread(target=worker, args=(task_queue, results_queue))
        # TODO: can't make the threads stop properly when an event is fired.
        # So for now we'll do it with daemon threads, which will get blocked
        # after the queue gets empty.
        t.daemon = True
        t.start()
    task_queue.join()
    return misc.unique(misc.flatten_list(results_queue.queue))

