import pymongo, re, datetime, json

from ivulncheck_api.lib import Vulns
from ivulncheck_api.lib.config_api import Configuration as conf
from ivulncheck_api.lib import functions as f

import logging
logging.basicConfig(filename=conf.getLogfile(), level=logging.INFO)

class DatabaseLayer(object):
    
    def init(self):
        
        self.vulndb = conf.getVulnDB()
        self.hostdb = conf.getHostDB()
        self.fixdb = conf.getFixDB()

    def init_test(self):
        self.testdb = conf.getTestDB()
        # Aliases to make tests work
        self.vulndb = conf.getTestDB()
        self.hostdb = conf.getTestDB()
        self.fixdb = conf.getTestDB()
        
    def reset_test(self):
        self.testdb.drop_collection('hosts')
        self.testdb.drop_collection('packages')
        self.testdb.drop_collection('vulns')
        self.testdb.drop_collection('fixes')
        self.testdb.drop_collection('pkg2cpe')
        self.testdb.drop_collection('debsecbugs')
        self.testdb.drop_collection('certfrs')


    # Functions
    # Function for converting data from database to python dict
    def sanitize(self, x):
      if type(x)==pymongo.cursor.Cursor:
        x=list(x)
      if type(x)==list:
        for y in x: self.sanitize(y)
      if x and  "_id" in x: x.pop("_id")
      return x
            
    # Index for CPEID
    def insertPkg2CPE(self, pkg2cpe):
        try:
            result = self.vulndb['pkg2cpe'].insert_one(pkg2cpe)
            if result.acknowledged == True:
                return True
            else:
                return False
        except pymongo.errors.DuplicateKeyError:
            return False
        
    def getPkg2CPE(self, pkgid):
        pkg2cpe = self.vulndb['pkg2cpe'].find_one({'id': pkgid})
        if pkg2cpe is not None:
            return pkg2cpe['cpeid']
        else:
            return None
            
    # MyDB Functions
        
    def getVulns(self, limit=False, query=[], skip=0, vulns=None):
        if type(query) == dict: query=[query]
        if type(vulns) == list: query.append({"id": {"$in": vulns}})
        if len(query) == 0:
            vuln=self.vulndb["vulns"].find().sort("bin_package", 1).limit(limit).skip(skip)
        elif len(query)  == 1:
            vuln=self.vulndb["vulns"].find(query[0]).sort("bin_package", 1).limit(limit).skip(skip)
        else:
            vuln=self.vulndb["vulns"].find({"$and": query}).sort("bin_package", 1).limit(limit).skip(skip)
        return self.sanitize(vuln)
      
    def getVulnsStats(self):
        nb_unknown = self.vulndb['vulns'].find({'urgency': ""}).count()
        nb_negligible = self.vulndb['vulns'].find({'urgency': "negligible"}).count()
        nb_low = self.vulndb['vulns'].find({'urgency': "low"}).count()
        nb_medium = self.vulndb['vulns'].find({'urgency': "medium"}).count()
        nb_high = self.vulndb['vulns'].find({'urgency': "high"}).count()
        nb_critical = self.vulndb['vulns'].find({'urgency': "critical"}).count()
        
        nb_rem_exp = self.vulndb['vulns'].find({'remotely_exp': True}).count()
        nb_non_rem_exp = self.vulndb['vulns'].find({'remotely_exp': False}).count()
        nb_rem_unknown = self.vulndb['vulns'].find({'remotely_exp': None}).count()
        
        stats = {'nb_unknown': nb_unknown, 'nb_negligible': nb_negligible, 'nb_low': nb_low, 'nb_medium': nb_medium, 'nb_high': nb_high, 'nb_critical':nb_critical, 'nb_rem_exp': nb_rem_exp, 'nb_non_rem_exp': nb_non_rem_exp, 'nb_rem_unknown': nb_rem_unknown}
        return stats
        
    def getPackages(self, limit=False, query=[], skip=0, hide_hosts=False):
        if type(query) == dict: query=[query]
        if len(query) == 0:
            pkg=self.hostdb["packages"].find().sort("id", 1).limit(limit).skip(skip)
        elif len(query)  == 1:
            if hide_hosts == True:
                pkg=self.hostdb["packages"].find(query[0], {'hosts': 0}).sort("id", 1).limit(limit).skip(skip)
            else:
                pkg=self.hostdb["packages"].find(query[0]).sort("id", 1).limit(limit).skip(skip)
        else:
            if hide_hosts == True:
                pkg=self.hostdb["packages"].find({"$and": query}, {'hosts': 0}).sort("id", 1).limit(limit).skip(skip)
            else:
                pkg=self.hostdb["packages"].find({"$and": query}).sort("id", 1).limit(limit).skip(skip)
        return self.sanitize(pkg)
      
    def getProducts(self, limit=False, query=[], skip=0):
      if type(query) == dict: query=[query]
      if len(query) == 0:
        pkg=self.hostdb["products"].find().sort("id", 1).limit(limit).skip(skip)
      elif len(query)  == 1:
        pkg=self.hostdb["products"].find(query[0]).sort("id", 1).limit(limit).skip(skip)
      else:
        pkg=self.hostdb["products"].find({"$and": query}).sort("id", 1).limit(limit).skip(skip)
      return self.sanitize(pkg)
    
    def getOSVersions(self):
        return self.hostdb["hosts"].distinct('os_version')
    
    def getSuites(self):
        return self.hostdb["packages"].distinct('os_version', filter={'os_type': "Debian"})
        
    def getReleases(self):
        return self.hostdb["packages"].distinct('os_version', filter={'os_type': "Ubuntu"})
         
    def getHosts(self, limit=False, query=[], skip=0):
      if type(query) == dict: query=[query]
      if len(query) == 0:
        hosts=self.hostdb["hosts"].find().sort("id", 1).limit(limit).skip(skip)
      elif len(query)  == 1:
        hosts=self.hostdb["hosts"].find(query[0]).sort("id", 1).limit(limit).skip(skip)
      else:
        hosts=self.hostdb["hosts"].find({"$and": query}).sort("id", 1).limit(limit).skip(skip)
      return self.sanitize(hosts)
      
    def getFixableHosts(self):
        return self.fixdb["fixes"].distinct('hosts')
        
    def getHostsStats(self):
        totvuln = self.vulndb["vulns"].count()
        totfix = self.fixdb["fixes"].count()
        
        hostnames = self.getFixableHosts()
        
        rules = json.load(open(conf.getRulesFile()))
        
        names_stats = f.hostnames2stats(hostnames, rules)
        
        # To escape (again) backslashes
        for name_stats in names_stats:
            name_stats['regex'].replace(r'\\', r'\\\\')
        
        stats = {'tot_vulns': totvuln, 'tot_fixes': totfix, 'names_stats': names_stats}
        
        return stats
        
    def getHostPackages(self, hostname):
        return self.sanitize(self.hostdb["packages"].find({'hosts': hostname}))
    
    def getHostStats(self, hostname):
        totpkg = self.hostdb["packages"].find({'hosts': hostname}).count()
        fixedpkg = self.fixdb["fixes"].find({'hosts': hostname}).count()
        nbvulnpkg = 0
        nbvuln = 0
        try:
            nbvuln = self.hostdb["hoststats"].find_one({'id': hostname})['nbvuln']
            nbvulnpkg = self.hostdb["hoststats"].find_one({'id': hostname})['nbvulnpkg']
        except Exception as e:
            logging.info('[%s] (%s)\tERROR : %s.', datetime.datetime.today().isoformat(), __name__, str(type(e)))
            # If no stats in database, calculates stats
            pkgs = self.hostdb["packages"].find({'hosts': hostname })
            for pkg in pkgs:
                nbpkgvulns = self.vulndb["vulns"].find({'bin_package': pkg['id'], 'bin_versions': pkg['version']}).count()
                nbvuln += nbpkgvulns
                if nbpkgvulns > 0:
                    nbvulnpkg += 1
            hoststats = {'id': hostname, 'nbvuln': nbvuln, 'nbvulnpkg': nbvulnpkg}
            result = self.insertHostStats(hoststats)
            if result == True:
                logging.info('[%s] (%s)\tHost Stats inserted for host : %s.', datetime.datetime.today().isoformat(), __name__, hostname)
        
        nbfix = self.fixdb["fixes"].find({'hosts': hostname}).count()
        
        stats = {'nb_vulns': nbvuln, 'nb_fixes': nbfix, 'up_packages' : totpkg - nbvulnpkg, 'fixed_packages' : fixedpkg, 'vuln_packages' : nbvulnpkg - fixedpkg}
        return stats
        
    def calculateHostStats(self, hostname):
        pkgs = self.hostdb["packages"].find({'hosts': hostname })
        nbvulnpkg = 0
        nbvuln = 0
        for pkg in pkgs:
            nbpkgvulns = self.vulndb["vulns"].find({'bin_package': pkg['id'], 'bin_versions': pkg['version']}).count()
            nbvuln += nbpkgvulns
            if nbpkgvulns > 0:
                nbvulnpkg += 1
        hoststats = {'id': hostname, 'nbvuln': nbvuln, 'nbvulnpkg': nbvulnpkg}
        return self.insertHostStats(hoststats)
        
    def getHostPackage(self, hostname, pkgid, pkgvers):
        return self.sanitize(self.hostdb["packages"].find_one({'hosts': hostname, 'id': pkgid, 'version': pkgvers}))
    
    def getHostVulns(self, hostname, limit=False, query=[], skip=0):
        # Keep queries
        orig_query=query[:]
        hostvulns = []
        hostpkgs = self.sanitize(self.hostdb["packages"].find())
        for hostpkg in hostpkgs:
            query.append({'bin_package': hostpkg['id']})
            query.append({'bin_versions': hostpkg['version']})
            vulns = self.getVulns(query=query)
            for vuln in vulns:
                hostvulns.append(vuln)
            # Restore query without bin_package and bin_versions
            query = orig_query[:]
        
        if limit == False:
            limit = len(hostvulns)
            
        return hostvulns[skip:skip+limit]
                
                      
    def getLastItem(self, collection):
        if self.vulndb[collection].count() > 0:
            return self.sanitize(self.vulndb[collection].find().limit(1).skip(self.vulndb[collection].count()-1))
        else:
            return None
            
    #~ def getFixes(self):
        #~ return self.sanitize(self.fixdb['fixes'].find())
        
    def getFixes(self, limit=False, query=[], skip=0, fixes=None, hide_hosts=False):
        if type(query) == dict: query=[query]
        if type(fixes) == list: query.append({"id": {"$in": fixes}})
        if len(query) == 0:
            fix=self.fixdb["fixes"].find().sort("src_package", 1).limit(limit).skip(skip)
        elif len(query)  == 1:
            if hide_hosts == True:
                fix=self.fixdb["fixes"].find(query[0], {'hosts': 0}).sort("src_package", 1).limit(limit).skip(skip)
            else:
                fix=self.fixdb["fixes"].find(query[0]).sort("src_package", 1).limit(limit).skip(skip)
        else:
            if hide_hosts == True:
                fix=self.fixdb["fixes"].find({"$and": query}, {'hosts': 0}).sort("src_package", 1).limit(limit).skip(skip)
            else:
                fix=self.fixdb["fixes"].find({"$and": query}).sort("src_package", 1).limit(limit).skip(skip)
        return self.sanitize(fix)
        
    def getFix(self, pkg_id, pkg_vers):
        return self.sanitize(self.fixdb['fixes'].find_one({'bin_package': pkg_id, 'bin_version': pkg_vers}))
        
    def getFixesStats(self):
        nb_unknown = self.fixdb['fixes'].find({'urgency': ""}).count()
        nb_negligible = self.fixdb['fixes'].find({'urgency': "negligible"}).count()
        nb_low = self.fixdb['fixes'].find({'urgency': "low"}).count()
        nb_medium = self.fixdb['fixes'].find({'urgency': "medium"}).count()
        nb_high = self.fixdb['fixes'].find({'urgency': "high"}).count()
        nb_critical = self.fixdb['fixes'].find({'urgency': "critical"}).count()
        
        nb_rem_exp = self.fixdb['fixes'].find({'remotely_exp': True}).count()
        nb_non_rem_exp = self.fixdb['fixes'].find({'remotely_exp': False}).count()
        nb_rem_unknown = self.fixdb['fixes'].find({'remotely_exp': None}).count()
        
        stats = {'nb_unknown': nb_unknown, 'nb_negligible': nb_negligible, 'nb_low': nb_low, 'nb_medium': nb_medium, 'nb_high': nb_high, 'nb_critical':nb_critical, 'nb_rem_exp': nb_rem_exp, 'nb_non_rem_exp': nb_non_rem_exp, 'nb_rem_unknown': nb_rem_unknown}
        return stats
            
    def getHostFixes(self, hostname):
        return self.sanitize(self.fixdb['fixes'].find({'hosts': hostname}))
        
    def getPackageFixes(self, pkg_src_id):
        return self.sanitize(self.fixdb['fixes'].find({'src_package': pkg_src_id}))
        
    def fixExists(self, hostname, pkg_id):
        if self.fixdb['fixes'].find({'hosts': hostname, 'package': pkg_id}).count() > 0:
            return True
        else:
            return False
            
    def skipFix(self, bin_package, bin_version):
        result = self.fixdb['fixes'].update_one({'bin_package': bin_package, 'bin_version': bin_version},{'$set': {'skipped': True}})
        if result.acknowledged == True:
            return True
        else:
            return False
        
    def unskipFix(self, bin_package, bin_version):    
        result = self.fixdb['fixes'].update_one({'bin_package': bin_package, 'bin_version': bin_version},{'$set': {'skipped': False}})
        if result.acknowledged == True:
            return True
        else:
            return False
            
    def insertHostStats(self, hoststats):
        try:
            result = self.hostdb['hoststats'].insert_one(hoststats)
            if result.acknowledged == True:
                return True
            else:
                return False
        except pymongo.errors.DuplicateKeyError:
            self.updateHostStats(hoststats)
            
    def updateHostStats(self, hoststats):
        # Get old hoststats
        old_hoststats = self.hostdb['hoststats'].find_one({'id': hoststats['id']})
        # Set new values
        old_hoststats['nbvuln'] = hoststats['nbvuln']
        old_hoststats['nbvulnpkg'] = hoststats['nbvulnpkg']
        result = self.hostdb['hoststats'].update_one({'id': hoststats['id']}, {"$set": old_hoststats})
        if result.acknowledged == True:
            return True
        else:
            return False

    def insertVuln(self, vuln):
        try:
            result = self.vulndb['vulns'].insert_one(vuln)
            if result.acknowledged == True:
                return True
            else:
                return False
        except pymongo.errors.DuplicateKeyError:
            self.updateVuln(vuln)

    def updateVuln(self, vuln):
        # Get old vuln
        old_vuln =  self.getVulns(query={'id': vuln['id'], 'bin_package': vuln['bin_package']})[0]
        # If source differ
        if old_vuln['src'] != vuln['src']:
            old_vuln['src'] = "both"
        # Add new versions
        if len(vuln['bin_versions']) > 0:
            for bin_vers in vuln['bin_versions']:
                # If versions differ keep old product
                if bin_vers not in old_vuln['bin_versions']:
                    old_vuln['bin_versions'].append(bin_vers)
        # Add new src_versions
        if len(vuln['src_versions']) > 0:
            for src_vers in vuln['src_versions']:
                # If versions differ keep old product
                if src_vers not in old_vuln['src_versions']:
                    old_vuln['src_versions'].append(src_vers)
        # Add new vuln_versions
        if len(vuln['vuln_versions']) > 0:
            for vuln_vers in vuln['vuln_versions']:
                # If versions differ keep old product
                if vuln_vers not in old_vuln['vuln_versions']:
                    old_vuln['vuln_versions'].append(vuln_vers)
        # Add new fixed_versions
        if len(vuln['fixed_versions']) > 0:
            for fix_vers in vuln['fixed_versions']:
                # If versions differ keep old product
                if fix_vers not in old_vuln['fixed_versions']:
                    old_vuln['fixed_versions'].append(fix_vers)
        # Update fixed tag
        if len(vuln['fixed']) > 0:
            if vuln['fixed'][0] not in old_vuln['fixed']:
                old_vuln['fixed'].append(vuln['fixed'][0])
        # Add published date
        if vuln['published_date'] != "":
            old_vuln['published_date'] = vuln['published_date']
                
        
        result = self.vulndb['vulns'].update_one({'id': old_vuln['id'], 'bin_package': old_vuln['bin_package']}, {"$set": old_vuln})
        if result.acknowledged == True:
            return True
        else:
            return False

    #~ def deleteVuln(self, vuln_id):
        #~ return self.deleteDocument("vulns", "id", vuln_id)
      
    def insertFix(self, fix):
        try:
            result = self.fixdb['fixes'].insert_one(fix)
            if result.acknowledged == True:
                return True
            else:
                return False
        except pymongo.errors.DuplicateKeyError:
            self.updateFix(fix)
        
    def updateFix(self, fix):
        # Get old fix
        old_fix =  self.getFixes(query={'bin_package': fix['bin_package'], 'bin_version': fix['bin_version']})[0]
        # If source differ
        if fix['src'] != old_fix['src']:
            old_fix['src'] = 'both'
        # get vuln number name as hits
        try:
            hits = old_fix['hits']
        except KeyError:
            hits = 1
            
        # Add new hosts
        if len(fix['hosts']) > 0:
            for host in fix['hosts']:
                if host not in old_fix['hosts']:
                    old_fix['hosts'].append(host)
        
        # Add new CVEID
        if len(fix['cve_ids']) > 0:
            for cve_id in fix['cve_ids']:
                # If CVEIDs differ add new CVEID
                if cve_id not in old_fix['cve_ids']:
                    old_fix['cve_ids'].append(cve_id)
        # Get max urgency
        if fix['urgency'] == "high":
            old_fix['urgency'] = "high"
        elif fix['urgency'] == "medium" and (old_fix['urgency'] == "low" or old_fix['urgency'] == ""):
            old_fix['urgency'] = "medium"
        elif fix['urgency'] == "low" and old_fix['urgency'] == "":
            old_fix['urgency'] = "low"
        # Add CVE published date
        if fix['cve_published_date'] != "":
            old_fix['cve_published_date'] = fix['cve_published_date']
        # Add Fix published date
        if fix['fix_published_date'] != "":
            old_fix['fix_published_date'] = fix['fix_published_date']
        
        result = self.fixdb['fixes'].update_one({'bin_package': old_fix['bin_package'], 'bin_version': fix['bin_version']}, {"$set": old_fix})
        if result.acknowledged == True:
            return True
        else:
            return False

    def insertPackage(self, package):
        try:
            result = self.hostdb['packages'].insert_one(package)
            if result.acknowledged == True:
                return True
            else:
                return False
        except pymongo.errors.DuplicateKeyError:
            self.updatePackage(package)
            
    def updatePackage(self, package):
        old_package = self.getPackages(query={'id': package['id'],'version': package['version']})[0]
        # Keep hosts
        for host in package['hosts']:
            if host not in old_package['hosts']:
                old_package['hosts'].append(host)
        result = self.hostdb['packages'].update_one({'id': old_package['id'], 'version': old_package['version']}, {'$set':  old_package})
        if result.acknowledged == True:
            return True
        else:
            return False
            
    def insertProduct(self, product):
        try:
            result = self.hostdb['products'].insert_one(product)
            if result.acknowledged == True:
                return True
            else:
                return False
        except pymongo.errors.DuplicateKeyError:
            self.updateProduct(product)
            
    def updateProduct(self, product):
        old_product = self.getProducts(query={'id': product['id'],'version': product['version']})[0]
        # Keep hosts
        for host in product['hosts']:
            if host not in old_product['hosts']:
                old_product['hosts'].append(host)
        result = self.hostdb['products'].update_one({'id': old_product['id'], 'version': old_product['version']}, {'$set':  old_product})
        if result.acknowledged == True:
            return True
        else:
            return False
            
    def updateUrgRmtFix(self, bin_package, bin_version, remotely_exp, urgency):
        result = self.fixdb['fixes'].update_one({'bin_package': bin_package, 'bin_version': bin_version}, {'$set': {'remotely_exp': remotely_exp, 'urgency': urgency}})
        if result.acknowledged == True:
            return True
        else:
            return False 
            
    def updateCPEPackage(self, id, version, cpe_id, cpe_src_id):
        result = self.hostdb['packages'].update_one({'id': id, 'version': version}, {'$set': {'cpe_id': cpe_id, 'cpe_src_id': cpe_src_id}})
        if result.acknowledged == True:
            return True
        else:
            return False
            
    def deletePackage(self, pkg_id, pkg_vers):
        return self.hostdb["packages"].delete_one({'id': pkg_id, 'version': pkg_vers})
        
    def insertHost(self, host):
        try:
            result = self.hostdb['hosts'].insert_one(host)
            if result.acknowledged == True:
                return True
            else:
                return False
        except pymongo.errors.DuplicateKeyError:
            return False
            
    def purgeHost(self, hostname):
        results = []
        if len(self.getHosts(query={'id': hostname})) > 0:
            results.append(self.purgeCollection("hosts", query={'id': hostname}))
            hosts = []
            hosts.append(hostname)
            results.append(self.purgeCollection("packages", query={'hosts': hosts}))
            results.append(self.purgeCollection("fixes", query={'hosts': hostname}))
        return results
            
      
    def purgeCollection(self, collection, query={'id': {'$ne': 'aaa'}}):
        if collection == "hosts":
            return self.hostdb[collection].delete_many(query)
        elif collection == "packages":
            return self.hostdb[collection].delete_many(query)
        elif collection == "products":
            return self.hostdb[collection].delete_many(query)
        elif collection == "vulns":
            return self.vulndb[collection].delete_many(query)
        elif collection == "pkg2cpe":
            return self.vulndb[collection].delete_many(query)
        elif collection == "fixes":
            return self.fixdb[collection].delete_many(query)
        
    def purgeAll(self):
        results = []
        results.append(self.purgeCollection("hosts"))
        results.append(self.purgeCollection("packages"))
        results.append(self.purgeCollection("products"))
        results.append(self.purgeCollection("vulns"))
        results.append(self.purgeCollection("pkg2cpe"))
        results.append(self.purgeCollection("fixes"))
        return results
