#!/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) 2018-2019  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 argparse
import sqlite3

# load the srt.sqlite schema indexes
dir_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
sys.path.insert(0, dir_path)
from common.srt_schema import ORM

# Setup:
verbose = False
cmd_skip = 0
cmd_count = 0

srtDbName = 'srt.sqlite'

#################################
# Common routines
#

# quick development/debugging support
def _log(msg):
    DBG_LVL =  os.environ['SRTDBG_LVL'] if ('SRTDBG_LVL' in os.environ) else 2
    DBG_LOG =  os.environ['SRTDBG_LOG'] if ('SRTDBG_LOG' in os.environ) else '/tmp/srt_dbg.log'
    if 1 == DBG_LVL:
        print(msg)
    elif 2 == DBG_LVL:
        f1=open(DBG_LOG, 'a')
        f1.write("|" + msg + "|\n" )
        f1.close()

#################################
# 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):
    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[ORM.DATASOURCE_ID]))
            is_change = True
        elif 'reset' == cmnd:
            commit_to_source(conn,(False,ds[ORM.DATASOURCE_ID]))
            is_change = True
        elif 'reset_not_nist' == cmnd:
            if 'nist' != ds[ORM.DATASOURCE_SOURCE]:
                print("RESETTING Data source [%s] data='%s' of '%s' load state from '%s' is '%s'" % (ds[ORM.DATASOURCE_ID],ds[ORM.DATASOURCE_DATA],ds[ORM.DATASOURCE_DESCRIPTION],ds[ORM.DATASOURCE_SOURCE],ds[ORM.DATASOURCE_LOADED]))
                commit_to_source(conn,(False,ds[ORM.DATASOURCE_ID]))
            else:
                commit_to_source(conn,(True,ds[ORM.DATASOURCE_ID]))
            is_change = True
        elif 'triage_keywords' == cmnd:
            if 'triage_keywords' == ds[ORM.DATASOURCE_DATA]:
                print("RESETTING Data source [%s] data='%s' of '%s' load state from '%s' is '%s'" % (ds[ORM.DATASOURCE_ID],ds[ORM.DATASOURCE_DATA],ds[ORM.DATASOURCE_DESCRIPTION],ds[ORM.DATASOURCE_SOURCE],ds[ORM.DATASOURCE_LOADED]))
                commit_to_source(conn,(False,ds[ORM.DATASOURCE_ID]))
                is_change = True
        else:
            print("Data source [%s] data='%s' of '%s' load state from '%s' is '%s'" % (ds[ORM.DATASOURCE_ID],ds[ORM.DATASOURCE_DATA],ds[ORM.DATASOURCE_DESCRIPTION],ds[ORM.DATASOURCE_SOURCE],ds[ORM.DATASOURCE_LOADED]))

    if is_change:
        conn.commit()


def settings():
    conn = sqlite3.connect(srtDbName)
    c = conn.cursor()

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

    for setting in c:
        print("Setting[%s] = '%s'" % (setting[ORM.SRTSETTING_NAME], setting[ORM.SRTSETTING_VALUE][0:40]))


#################################
# remove_app_sources
#

def remove_app_sources(master_app):
    conn = sqlite3.connect(srtDbName)
    cur = conn.cursor()
    cur_write = conn.cursor()

    # Scan the SRTool Settings
    cur.execute("SELECT * FROM orm_datasource")

    is_change = False
    for setting in cur:
        if master_app == setting[ORM.DATASOURCE_SOURCE]:
            print("Deleting [%s] = '%s','%s'" % (setting[ORM.DATASOURCE_ID], setting[ORM.DATASOURCE_SOURCE], setting[ORM.SRTSETTING_VALUE]))
            sql = 'DELETE FROM orm_datasource WHERE id=?'
            cur_write.execute(sql, (setting[ORM.DATASOURCE_ID],))
            is_change = True

    if is_change:
        conn.commit()
    conn.close()

#################################
# fix_new_reserved
#

# Is this reserved by Mitre? Is '** RESERVED **' within the first 20 char positions?
def fix_new_reserved():
    global cmd_skip
    global cmd_count

    conn = sqlite3.connect(srtDbName)
    cur = conn.cursor()
    cur_write = conn.cursor()

    cur.execute('SELECT * FROM orm_cve WHERE status = "%s"' % ORM.STATUS_NEW)
    i = 0
    j = 0
    for cve in cur:
        i += 1

        # Progress indicator support
        if 0 == i % 10:
            print('%05d: %20s\r' % (i,cve[ORM.CVE_NAME]), end='')
        if (0 == i % 200):
            conn.commit()
        # Development/debug support
        if cmd_skip:
            if i < cmd_skip:
                continue
        if cmd_count:
            if (i - cmd_skip) > cmd_count:
                print("Count return: %s,%s" % (i,cmd_count))
                break

        reserved_pos = cve[ORM.CVE_DESCRIPTION].find('** RESERVED **')
        if (0 <= reserved_pos) and (20 > reserved_pos):
            print("STATUS_NEW_RESERVED:%s:%s:%s" % (cve[ORM.CVE_STATUS],cve[ORM.CVE_NAME],cve[ORM.CVE_DESCRIPTION][:40]))
            sql = ''' UPDATE orm_cve
                      SET status = ?
                      WHERE id = ?'''
            cur_write.execute(sql, (ORM.STATUS_NEW_RESERVED, cve[ORM.CVE_ID],))
            j += 1
    print("\nCVE COUNT=%5d,%5d" % (i,j))
    conn.commit()

#################################
# main loop
#
def main(argv):
    global verbose
    global cmd_skip
    global cmd_count

    # setup
    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('--settings', '-S', action='store_const', const='settings', dest='command', help='Show the SRT Settings')
    parser.add_argument('--remove-app-sources', dest='remove_app_sources', help='Remove data sources for a previous app')

    parser.add_argument('--force', '-f', action='store_true', dest='force', help='Force the update')
    parser.add_argument('--verbose', '-v', action='store_true', dest='verbose', help='Debugging: verbose output')
    parser.add_argument('--skip', dest='skip', help='Debugging: skip record count')
    parser.add_argument('--count', dest='count', help='Debugging: short run record count')

    parser.add_argument('--fix-new-reserved', action='store_const', const='fix_new_reserved', dest='command', help='Reset new reserved CVEs to NEW_RESERVED')

    args = parser.parse_args()

    master_log = open(os.path.join(script_pathname, "update_logs/master_log.txt"), "a")

    verbose = args.verbose
    if None != args.skip:
        cmd_skip = int(args.skip)
    if None != args.count:
        cmd_count = int(args.count)

    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 'reset_sources' == args.command:
        sources('reset')
    elif 'settings' == args.command:
        settings()
    elif 'fix_new_reserved' == args.command:
        fix_new_reserved()
    elif args.remove_app_sources:
        remove_app_sources(args.remove_app_sources)
    else:
        print("Command not found")
    master_log.close()

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