path: root/bin/srtool.py

#!/usr/bin/env python3
#
# ex:ts=4:sw=4:sts=4:et
# -*- tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*-
#
# Security Response Tool Commandline Tool
#
# Copyright (C) 2017       Wind River Systems
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

import os
import sys
import re
import csv
import xml.etree.ElementTree as ET
import argparse
import sqlite3
import subprocess
import json

# setup
lookupTable = []
cveIndex = {}
jiraIndex = {}
db_change = False

srtDbName = 'srt.sqlite'

is_verbose = False


#################################
# reset sources
#

# source_data = (source_id)
def commit_to_source(conn, source_data):
    sql = ''' UPDATE orm_datasource
              SET loaded = ?
              WHERE id = ?'''
    cur = conn.cursor()
    print("UPDATE_SCORE:%s" % str(source_data))
    cur.execute(sql, source_data)

def sources(cmnd):

    DS_ID = 0
    DS_DATA = 1
    DS_SOURCE = 2
    DS_TYPE = 3
    DS_DESCRIPTION = 4
    DS_FILE_PATH = 5
    DS_URL = 6
    DS_LOADED = 7
    
    conn = sqlite3.connect(srtDbName)
    c = conn.cursor()

    print('Sources(%s)' % cmnd)

    c.execute("SELECT * FROM orm_datasource")
    is_change = False
    for ds in c:
        if 'set' == cmnd:
            commit_to_source(conn,(True,ds[DS_ID]))
            is_change = True
        elif 'reset' == cmnd:
            commit_to_source(conn,(False,ds[DS_ID]))
            is_change = True
        elif 'reset_not_nist' == cmnd:
            if 'nist' != ds[DS_SOURCE]:
                print("RESETTING Data source [%s] data='%s' of '%s' load state from '%s' is '%s'" % (ds[DS_ID],ds[DS_DATA],ds[DS_DESCRIPTION],ds[DS_SOURCE],ds[DS_LOADED]))
                commit_to_source(conn,(False,ds[DS_ID]))
            else:
                commit_to_source(conn,(True,ds[DS_ID]))
            is_change = True
        elif 'triage_keywords' == cmnd:
            if 'triage_keywords' == ds[DS_DATA]:
                print("RESETTING Data source [%s] data='%s' of '%s' load state from '%s' is '%s'" % (ds[DS_ID],ds[DS_DATA],ds[DS_DESCRIPTION],ds[DS_SOURCE],ds[DS_LOADED]))
                commit_to_source(conn,(False,ds[DS_ID]))
                is_change = True
        else:
            print("Data source [%s] data='%s' of '%s' load state from '%s' is '%s'" % (ds[DS_ID],ds[DS_DATA],ds[DS_DESCRIPTION],ds[DS_SOURCE],ds[DS_LOADED]))

    if is_change:
        conn.commit()


#################################
# update_scores
#

keywords_for = []
keywords_against = []

csvfile_name = 'data/keyword_filters_full.csv'

def read_keywords(csvfile_name):
    # mode,type,keyword,weight
    # y,key,abiword,

    global keywords_for
    global keywords_against
    
    KEY_MODE=0
    KEY_TYPE=1
    KEY_KEY=2
    KEY_WEIGHT=3

    i_index=0
    is_header = True
    with open(csvfile_name, newline='') as csvfile:
        CPE_reader = csv.reader(csvfile, delimiter=',', quotechar='"')
        for row in CPE_reader:
            if is_header or not len(row):
                is_header = False
                continue

            if (KEY_WEIGHT+1) != len(row):
                print("KEY_ROWLEN_ERROR:'%s'" % row)
                continue

            i_index += 1
            if 0 == i_index % 100:
                print('%04d: %20s\r' % (i_index,row[KEY_KEY]), end='', file=sys.stderr)

#                # DEBUG ### TODO
#                if 0 < Command.debug_jira_limit:
#                    if i_index > Command.debug_jira_limit:
#                        return

            key = row[KEY_MODE]
            if '#' == key[0]:
                key = key[1:]

            if 'y' == key:
                keywords_for.append("%s,%s" % (row[KEY_KEY].lower(),row[KEY_WEIGHT]))
            elif 'n' == key:
                keywords_against.append("%s,%s" % (row[KEY_KEY].lower(),row[KEY_WEIGHT]))

    print('keywords_for = %d' % len(keywords_for))
    print('keywords_against = %d' % len(keywords_against))


# score_data = (score,cve_id)
def commit_score(conn, score_data):
    sql = ''' UPDATE orm_investigation
              SET score = ?
              WHERE id = ?'''
    cur = conn.cursor()
#    print("UPDATE_SCORE:%s" % str(score_data))
    cur.execute(sql, score_data)

def cve_score(is_test):
    global is_verbose    
    global csvfile_name
    global keywords_for
    global keywords_against
    
    read_keywords(csvfile_name)

    conn = sqlite3.connect(srtDbName)
    c = conn.cursor()

    CVE_ID = 0
    CVE_NAME = 1
    CVE_SOURCE = 2
    CVE_STATUS = 3
    CVE_WR_COMMENTS = 4
    CVE_WR_COMMENTS_PRIVATE = 5
    CVE_CVE_DATA_TYPE = 6
    CVE_CVE_DATA_FORMAT = 7
    CVE_CVE_DATA_VERSION = 8
    CVE_PUBLIC = 9
    CVE_PUBLISH = 10
    CVE_PUBLISH_DATE = 11
    CVE_DESCRIPTION = 12
    CVE_PUBLISHEDDATE = 13
    CVE_LASTMODIFIEDDATE = 14
    CVE_RECOMMEND = 15
    CVE_CPE_LIST = 16

    # Scan the CVEs
    c.execute("SELECT * FROM orm_cve")

    index = 0
    count = 0
    is_change = False
    for cve in c:

        list_for = ''
        list_against = ''
        total = 0

        for keypair in keywords_for:
            #print("keypair='%s'" % keypair)
            key,w = keypair.split(',')
            weight = 1
            if w:
                weight = int(w)
# re.search(r'\bis\b', your_string)
            if ' '+key+' ' in ' '+cve[CVE_DESCRIPTION].lower()+' ':
                list_for += "%s," % key
                total += weight

        desc = ' '+cve[CVE_DESCRIPTION].lower()+' '
        for keypair in keywords_against:
            #print("keypair='%s'" % keypair)
            key,w = keypair.split(',')
            weight = -1
            if w:
                weight = int(w)
#            if key in cve[CVE_DESCRIPTION].lower():
            if ' '+key+' ' in desc:
                list_against += "%s," % key
                total += weight

        if list_for or list_against:
#            print("CVE=%s, DESC=%s, Total=%d, Y=%s, N=%s" % (cve[CVE_NAME],cve[CVE_DESCRIPTION][:20],total,list_for[:-1],list_against[:-1]))
            print("%s\t%s\t%s\t%s\t%s\t%s" % (cve[CVE_NAME],cve[CVE_STATUS],total,list_for[:-1],list_against[:-1],cve[CVE_DESCRIPTION]))
            
        index += 1
        if 0 == index % 10:
            print("%4d : %30s \r" % (index,cve[CVE_NAME]),end='', file=sys.stderr)
            count += 1
#        if index > 200:
#            break

    if is_change:
        conn.commit()
    print("Done (%d of %d)" % (count,index))


# score_data = (score,cve_id)
def commit_score(conn, score_data):
    sql = ''' UPDATE orm_investigation
              SET score = ?
              WHERE id = ?'''
    cur = conn.cursor()
#    print("UPDATE_SCORE:%s" % str(score_data))
    cur.execute(sql, score_data)

def settings():
    global is_verbose    
    global csvfile_name
    global keywords_for
    global keywords_against
    
    read_keywords(csvfile_name)

    conn = sqlite3.connect(srtDbName)
    c = conn.cursor()

    SETTING_ID = 0
    SETTING_NAME = 1
    SETTING_HELP = 2
    SETTING_VALUE = 3

    # Scan the CVEs
    c.execute("SELECT * FROM orm_srtsetting")

    index = 0
    count = 0
    is_change = False
    for setting in c:
        print("Setting[%s] = '%s'" % (setting[SETTING_NAME], setting[SETTING_VALUE][0:40]))

    # Scan the CVEs
    c.execute("SELECT * FROM orm_srtsetting where name = '%s'" % 'keywords_for')
    
    setting = c.fetchone()
    print("Setting2[%s] = '%s'" % (setting[SETTING_NAME], setting[SETTING_VALUE].split('|')[0]))
    


#################################
# main loop
#

def main(argv):

    # setup
    is_test = True

    parser = argparse.ArgumentParser(description='srtool.py: manage the SRTool database')

    parser.add_argument('--sources', '-s', nargs='?', const='display', help='SRTool Sources')
    parser.add_argument('--reset-sources', '-r', action='store_const', const='reset_sources', dest='command', help='Reset SRTool Sources')
    parser.add_argument('--triage-scores', '-t', action='store_const', const='cve_score', dest='command', help='Score the CVEs')
    parser.add_argument('--settings', '-S', action='store_const', const='settings', dest='command', help='Show the SRT Settings')
    parser.add_argument('--test', '-T', action='store_true', dest='is_test', help='Test database against SRTFI')
    args = parser.parse_args()

#    print('Args = %s' % args)
#    return
    
    is_test = False
    if None != args.is_test:
        is_test = args.is_test

    if args.sources:
        if args.sources.startswith('s'):
            sources("set")
        elif 0 <= args.sources.find('nist'):
            sources("reset_not_nist")
        elif args.sources.startswith('r'):
            sources("reset")
        elif args.sources.startswith('t'):
            sources("triage_keywords")
        else:
            sources("display")
    elif 'cve_score' == args.command:
        cve_score(is_test)
    elif 'settings' == args.command:
        settings()
    else:
        print("Command not found")

if __name__ == '__main__':
    global script_pathname
    script_pathname=os.path.dirname(sys.argv[0])
    main(sys.argv[1:])