Commit 82efadcc authored by Alexandre Dulaunoy's avatar Alexandre Dulaunoy
Browse files

Merge pull request #99 from PidgeyL/master 

Initial work on Database Layer remodelling
parents 1b4294df 66c15453

bin/db_dump.py

+2 −12
Original line number Diff line number Diff line
@@ -17,19 +17,9 @@ import json 
from bson import json_util



import lib.CVEs as cves

import lib.DatabaseLayer as db

from lib.Config import Configuration



# connect to db

db = Configuration.getMongoConnection()

collection = db.cves





def dumpallcveid(limit=False):

    cveid = []

    for x in collection.find({}).sort('_id', 1).limit(limit):

        cveid.append(x['id'])

    return cveid



argParser = argparse.ArgumentParser(description='Dump database in JSON format')

argParser.add_argument('-r', default=False, action='store_true', help='Include ranking value')

argParser.add_argument('-v', default=False, action='store_true', help='Include vfeed map')
@@ -43,6 +33,6 @@ capeclookup = args.c 


l = cves.last(rankinglookup=rankinglookup, vfeedlookup=vfeedlookup, capeclookup=capeclookup)



for cveid in dumpallcveid(limit=args.l):

for cveid in db.getCVEIDs(limit=args.l):

    item = l.getcve(cveid=cveid)

    print (json.dumps(item, sort_keys=True, default=json_util.default))

bin/search_cpe.py

+8 −20
Original line number Diff line number Diff line
@@ -16,17 +16,14 @@ import re 
import argparse



from lib.Config import Configuration



import lib.DatabaseLayer as db



runPath = os.path.dirname(os.path.realpath(__file__))



# connect to DB

db = Configuration.getMongoConnection()



vOutput = ""



argParser = argparse.ArgumentParser(description='Search for CPE with a pattern')

argParser.add_argument('-s', type=str, help='search in cpe list')

argParser.add_argument('-s', type=str, required=True, help='search in cpe list')

argParser.add_argument('-o', type=str, help='O = output format [compact]')

argParser.add_argument('-f', action='store_true', help='Enlarge the CPE search to all CPE indexed. Need the cpeother activated.', default=False)
@@ -35,9 +32,8 @@ cpeSearch = args.s 
vOutput = args.o





def search(collection, cpe):

    res = collection.find({'id': {'$regex': re.compile(cpe, re.IGNORECASE)}})

    res.count()

def search(cpe):

    res = db.getCPEMatching(re.compile(cpe, re.IGNORECASE), args.f)



    if vOutput == "compact":

        for item in res:
@@ -46,16 +42,8 @@ def search(collection, cpe): 
        for item in res:

            print(item['id'] + "  " + item['title'])





if not cpeSearch:

    print ("no option provided")

    argParser.print_help()

else:

# replace special characters in cpeSearch with encoded version.

cpeSearch = re.sub(r'\(', '%28', cpeSearch)

cpeSearch = re.sub(r'\)', '%29', cpeSearch)



    search(db.cpe, cpeSearch)



    if args.f:

        search(db.cpeother, cpeSearch)
 
search(cpeSearch)

bin/search_xmpp.py

+0 −4
Original line number Diff line number Diff line
@@ -31,7 +31,6 @@ from optparse import OptionParser 
import sleekxmpp

import json



from lib.Config import Configuration

from lib.Query import lastentries, apigetcve, apibrowse, apisearch

# BSON MongoDB include ugly stuff that needs to be processed for standard JSON

from bson import json_util
@@ -46,9 +45,6 @@ runPath = os.path.dirname(os.path.realpath(__file__)) 


rankinglookup = True



db = Configuration.getMongoConnection()

collection = db.cves



helpmessage = "\nlast <n> cve entries (output: JSON) \n"

helpmessage = helpmessage + "cvetweet <n> cve entries (output: Text) \n"

helpmessage = helpmessage + "browse vendors and products (output: JSON)\n"

lib/CVEs.py

+3 −8
Original line number Diff line number Diff line
@@ -21,6 +21,7 @@ import math 


from lib.Config import Configuration

from lib.Toolkit import exploitabilityScore,impactScore

import lib.DatabaseLayer as db



class last:
@@ -40,9 +41,6 @@ class last: 


        if rankinglookup:

            self.ranking = connectdb['ranking']

        if namelookup:

            self.cpeOther = connectdb['cpeother']

            self.cpe = connectdb['cpe']

        if vfeedlookup:

            self.vfeed = connectdb['vfeed']

        if capeclookup:
@@ -67,14 +65,11 @@ class last: 
        return ref



    def getcpe(self, cpeid=None):



        if not(self.namelookup):

            return cpeid



        e = self.cpe.find_one({'id': cpeid})



        e = db.getCPE(cpeid)

        if e is None:

            e = self.cpeOther.find_one({'id': cpeid})

            e = db.getAlternativeCPE(cpeid)

            if e is None:

                return cpeid

        if 'id' in e:

lib/DatabaseLayer.py

0 → 100644
+172 −0
Original line number Diff line number Diff line 
#!/usr/bin/env python3.3

# -*- coding: utf-8 -*-

#

# Database layer

#  translates database calls to functions

#

# Software is free software released under the "Modified BSD license"

#



# Copyright (c) 2014-2015       Pieter-Jan Moreels - pieterjan.moreels@gmail.com



# imports

from lib.Config import Configuration as conf

import pymongo

import re



# Variables

db=conf.getMongoConnection()

colCVE=      db['cves']

colCPE=      db['cpe']

colCPEOTHER= db['cpeother']

colWHITELIST=db['mgmt_whitelist']

colBLACKLIST=db['mgmt_blacklist']

colUSERS=    db['mgmt_users']

colSEEN=     db['mgmt_seen']

colINFO=     db['info']

colVFEED=    db['vfeed']

# Functions

def sanitize(x):

  if type(x)==pymongo.cursor.Cursor:

    x=list(x)

  if type(x)==list:

    for y in x: sanitize(y)

  if x and  "_id" in x: x.pop("_id")

  return x



# DB Functions

def ensureIndex(col, field):

  db[col].ensure_index(field)



def setColUpdate(collection, date):

  colINFO.update({"db": collection}, {"$set": {"last-modified": date}}, upsert=True)



def insertCVE(cve):

  colCVE.insert(cve)



def updateCVE(cve):

  colCVE.update({"id": cve['id']}, {"$set": {"cvss": cve['cvss'], "summary": cve['summary'], "references": cve['references'],

                                             "cwe": cve['cwe'], "vulnerable_configuration": cve['vulnerable_configuration'],

                                             "vulnerable_configuration_cpe_2_2": cve['vulnerable_configuration_cpe_2_2'], 'last-modified': cve['Modified']}})



# API Functions

def cvesForCPE(cpe):

  if not cpe: return []

  return sanitize(colCVE.find({"vulnerable_configuration": {"$regex": cpe}}).sort("Modified", -1))



# User Functions

def seenCVEs(user):

  data = colSEEN.find_one({"user": user})

  if not data:

    colSEEN.insert({"user": user, "seen_cves": []})

    return []

  else:

    return data['seen_cves']



def addSeenCVEs(user, CVEs):

  if type(CVEs) == str: CVEs=[CVEs]

  if type(CVEs) == list:

    seen=list(set(CVEs)-set(seenCVEs(user)))

    if seen:

      colSEEN.update({"user": user},{"$addToSet": {"seen_cves": { "$each": seen}}})



def removeSeenCVEs(user, CVEs):

  if type(CVEs) == str: CVEs=[CVEs]

  if type(CVEs) == list:

    colSEEN.update({"user": user}, {"$pullAll": {"seen_cves": CVEs}})



def isMasterAccount(user):

  return False if colUSERS.find({"username": user, "master": True}).count() == 0 else True



def userExists(user):

  return True if colUSERS.find({"username": user}).count() > 0 else False



def isSingleMaster(user):

  return True if len(list(colUSERS.find({"username": {"$ne": user}, "master": True}))) == 0 else False



# Query Functions

# Generic data

def getCVEs(limit=-1, query=[], skip=0):

  if type(query) == dict: query=[query]

  if len(query) == 0:

    cves=colCVE.find().sort("Modified", -1).limit(limit).skip(skip)

  elif len(query)  == 1:

    cves=colCVE.find(query[0]).sort("Modified", -1).limit(limit).skip(skip)

  else:

    cves=colCVE.find({"$and": query}).sort("Modified", -1).limit(limit).skip(skip)

  return sanitize(cves)



def getCVEIDs(limit=-1):

  return [x["id"] for x in colCVE.find().limit(limit).sort("Modified", -1)]



def getCVE(id):

  return sanitize(colCVE.find_one({"id": id}))



def getCPE(id):

  return sanitize(colCPE.find_one({"id": id}))



def getCPEMatching(regex, fullSearch=False):

  lst=list(colCPE.find({"id": {"$regex": regex}}))

  if fullSearch: lst.extend(colCPEOTHER.find({"id": {"$regex": regex}}))

  return lst



def getAlternativeCPE(id):

  return sanitize(colCPEOTHER.find_one({"id": id}))



def getFreeText(text):

  return [x["obj"] for x in db.command("text", "cves", search=text)["results"]]



def getInfo(collection):

  return sanitize(colINFO.find_one({"db": collection}))



def getLastModified(collection):

  info=getInfo(collection)

  return info['last-modified'] if info else None



def getSize(collection):

  return db[collection].count()



def vFeedLinked(key, val):

  print(key)

  print(val)

  cveList=[x['id'] for x in colVFEED.find({key: val})]

  return sanitize(getCVEs(query={'id':{'$in':cveList}}))



def getDBStats():

  cols=['cve', 'cpe', 'cpeOther', 'capec', 'd2sec', 'vendor', 'vfeed']

  stats={x+'A': getSize(x.lower()) for x in cols}

  stats['cveA']=getSize('cves')

  stats.update({x+'U': getLastModified(x.lower()) for x in cols})

  stats.update({'blA': colBLACKLIST.count(), 'wlA':colWHITELIST.count()})

  stats.update({'dbOnDisk': db.command("dbstats")['storageSize'], 'dbSize':db.command('dbstats')['dataSize']})

  stats['dbName']=conf.getMongoDB()

  return stats



# Dynamic data

def getWhitelist():

  return sanitize(colWHITELIST.find())



def getRules(list):

  if list.lower()=='whitelist':

    col=colWHITELIST

  elif list.lower()=='blacklist':

    col=colBLACKLIST

  else:

    return []

  rlist=col.find({'type':'cpe'}).distinct('id')

  hardware=["cpe:2.3:([^:]*:){9}"+re.escape(x) for x in col.find({'type':'targethardware'}).distinct('id')]

  software=["cpe:2.3:([^:]*:){8}"+re.escape(x) for x in col.find({'type':'targetsoftware'}).distinct('id')]

  rlist.extend(hardware)

  rlist.extend(software)

  return rlist





def getBlacklist():

  return sanitize(colBLACKLIST.find())



# Users

def getUsers():

  return sanitize(colUSERS.find())



def getUser(user):

  return sanitize(colUSERS.find_one({"username": user}))

CRA Git | Maintained and supported by SUSTech CRA and CCSE