cve-kb_db.py

"""Simple python program making an API call to Microsoft Graphs to fetch a map of CVEs and KBs"""

from oauthlib.oauth2 import BackendApplicationClient
from requests_oauthlib import OAuth2Session

import apiconfig

import mysql.connector

def getCVE_KB():
	client = BackendApplicationClient(client_id = apiconfig.CLIENT_ID)

	oauth = OAuth2Session(client=client)

	token = oauth.fetch_token(token_url = apiconfig.AUTHORITY_URL + apiconfig.TOKEN_ENDPOINT,
		client_id = apiconfig.CLIENT_ID,
		client_secret = apiconfig.CLIENT_SECRET,
		scope = apiconfig.SCOPE)

	endpoint = apiconfig.RESOURCE + apiconfig.RESOURCE_ENDPOINT
	headers = {'Content-type': 'application/json',
		   'Bearer': token['access_token']}

	graphdata = oauth.get(endpoint, headers=headers).json()
	
	return graphdata

#print(graphdata['value'][0])

def connect_db(dbHost, dbUser, dbPass, dbName):

	db = mysql.connector.connect(
	  host = dbHost,
	  user = dbUser,
	  passwd = dbPass,
	  database = dbName
	)
	return db
	
def createTable_db():
	
	db = connect_db(apiconfig.dbHost,
			apiconfig.dbUser,
			apiconfig.dbPass,
			apiconfig.dbName
			)

	cursor = db.cursor()
	
	use_dbQuery = "USE " + apiconfig.dbName
	cursor.execute(use_dbQuery)
	
	createTable_dbQuery = "CREATE TABLE cve_kb (id INT NOT NULL AUTO_INCREMENT,\
						     cvekbId VARCHAR(255),\
						     cveID VARCHAR(255),\
						     kbId VARCHAR(255),\
						     title VARCHAR(255),\
						     severity VARCHAR(255),\
						     PRIMARY KEY (id)\
						     )"
	cursor.execute(createTable_dbQuery)

def insert_db():
	
	db = connect_db(apiconfig.dbHost,
			apiconfig.dbUser,
			apiconfig.dbPass,
			apiconfig.dbName
			)

	cursor = db.cursor()
	
	graphdata = getCVE_KB()
	
	n_graphdata = graphdata['@odata.count']
	for elem in range(n_graphdata):
		insert_dbQuery = "INSERT INTO cve_kb (cvekbId, cveID, kbId, title, severity)\
							VALUES (%s, %s, %s, %s, %s)"
		insert_dbVal = tuple(graphdata['value'][elem].values())
		cursor.execute(insert_dbQuery, insert_dbVal)
		db.commit()

createTable_db()
insert_db()