
from base.config import vwConfig
from frameworks.nessus import NessusAPI
from utils.cli import bcolors
import pandas as pd
import sys
import os
import io
import time
import sqlite3

# TODO Create logging option which stores data about scan
import logging



class vulnWhisperer(object):

    def __init__(self, config=None, db_name='report_tracker.db', purge=False, verbose=None, debug=False, username=None, password=None):

        self.verbose = verbose
        self.nessus_connect = False
        self.develop = True
        self.purge = purge


        if config is not None:
            try:
                self.config = vwConfig(config_in=config)
                self.nessus_enabled = self.config.getbool('nessus', 'enabled')

                if self.nessus_enabled:
                    self.nessus_hostname = self.config.get('nessus', 'hostname')
                    self.nessus_port = self.config.get('nessus', 'port')

                    if password:
                        self.nessus_password = password
                    else:
                        self.nessus_password = self.config.get('nessus', 'password')
                    

                    if username:
                        self.nessus_username = username
                    else:
                        self.nessus_username = self.config.get('nessus', 'username')
                    
                    self.nessus_writepath = self.config.get('nessus', 'write_path')
                    self.nessus_dbpath = self.config.get('nessus', 'db_path')
                    self.nessus_trash = self.config.getbool('nessus', 'trash')
                    self.verbose = self.config.getbool('nessus', 'verbose')

                    try:
                        self.vprint(
                            '{info} Attempting to connect to nessus...'.format(info=bcolors.INFO))
                        self.nessus = NessusAPI(hostname=self.nessus_hostname,
                                                port=self.nessus_port,
                                                username=self.nessus_username,
                                                password=self.nessus_password)
                        self.nessus_connect = True
                        self.vprint(
                            '{success} Connected to nessus on {host}:{port}'.format(success=bcolors.SUCCESS,
                                                                                                        host=self.nessus_hostname,
                                                                                                        port=str(self.nessus_port)))
                    except Exception as e:
                        self.vprint(e)
                        raise Exception(
                            "{fail} Could not connect to nessus -- Please verify your settings in {config} are correct and try again.\nReason: {e}".format(config=self.config,
                                                                                                                                                           fail=bcolors.FAIL,
                                                                                                                                                           e=e))

            except Exception as e:

                self.vprint('{fail} Could not properly load your config!\nReason: {e}'.format(fail=bcolors.FAIL, e=e))
                sys.exit(0)

        if db_name is not None:
            if self.nessus_dbpath:
                self.database = os.path.join(self.nessus_dbpath, db_name)
            else:
                self.database = os.path.abspath(os.path.join(os.path.dirname( __file__ ), 'database', db_name))

            try:
                self.conn = sqlite3.connect(self.database)
                self.cur = self.conn.cursor()
                self.vprint("{info} Connected to database at {loc}".format(info=bcolors.INFO, loc=self.database))
            except Exception as e:
                self.vprint("{fail} Could not connect to database at {loc}\nReason: {e} - Please ensure the path exist".format(e=e, fail=bcolors.FAIL, loc=self.database))

        else:
            self.vprint('{fail} Please specify a database to connect to!'.format(fail=bcolors.FAIL))
            exit(0)

        self.table_columns = ['scan_name',
                              'scan_id',
                              'last_modified',
                              'filename',
                              'download_time',
                              'record_count',
                              'source',
                              'uuid',
                              'processed']
        self.init()
        self.uuids = self.retrieve_uuids()
        self.processed = 0
        self.skipped = 0
        self.scan_list = []



    def vprint(self, msg):
        if self.verbose:
            print(msg)


    def create_table(self):
        self.cur.execute("create table if not exists scan_history (id integer primary key, scan_name text, scan_id integer, last_modified date, filename text, download_time date, record_count integer, source text, uuid text, processed integer)")
        self.conn.commit()

    def delete_table(self):
        self.cur.execute('drop table if exists scan_history')
        self.conn.commit()

    def init(self):
        if self.purge:
            self.delete_table()
        self.create_table()

    def cleanser(self, _data):
        repls = ('\n', '|||'), ('\r', '|||'), (',',';')
        data = reduce(lambda a, kv: a.replace(*kv), repls, _data)
        return data

    def path_check(self, _data):
        if self.nessus_writepath:
            data = self.nessus_writepath + '/' + _data
        return data

    def scan_count(self, scans, completed=False):
        """

        :param scans: Pulls in available scans
        :param completed: Only return completed scans
        :return:
        """
        self.vprint('{info} Gathering all scan data... this may take a while...'.format(info=bcolors.INFO))
        scan_records = []
        for s in scans:
            if s:
                record = {}
                record['scan_id'] = s['id']
                record['scan_name'] = s.get('name', '')
                record['owner'] = s.get('owner', '')
                record['creation_date'] = s.get('creation_date', '')
                record['starttime'] = s.get('starttime', '')
                record['timezone'] = s.get('timezone', '')
                record['folder_id'] = s.get('folder_id', '')
                try:
                    for h in self.nessus.get_scan_history(s['id']):
                        record['uuid'] = h.get('uuid', '')
                        record['status'] = h.get('status', '')
                        record['history_id'] = h.get('history_id', '')
                        record['last_modification_date'] = h.get('last_modification_date', '')
                        record['norm_time'] = self.nessus.get_utc_from_local(int(record['last_modification_date']),
                                                                   local_tz=self.nessus.tz_conv(record['timezone']))
                        scan_records.append(record.copy())


                except Exception as e:
                    #print(e)
                    pass

        if completed:
            scan_records = [s for s in scan_records if s['status'] == 'completed']
        return scan_records


    def record_insert(self, record):
        self.cur.execute("insert into scan_history({table_columns}) values (?,?,?,?,?,?,?,?,?)".format(
            table_columns=', '.join(self.table_columns)), record)
        self.conn.commit()

    def retrieve_uuids(self):
        """
        Retrieves UUIDs from database and checks list to determine which files need to be processed.
        :return:
        """
        self.conn.text_factory = str
        self.cur.execute('select uuid from scan_history')
        results = frozenset([r[0] for r in self.cur.fetchall()])
        return results


    def whisper_nessus(self):
        if self.nessus_connect:
            scan_data = self.nessus.get_scans()
            folders = scan_data['folders']
            scans = scan_data['scans']
            all_scans = self.scan_count(scans)
            if self.uuids:
                scan_list = [scan for scan in all_scans if (scan['uuid'] not in self.uuids and scan['status']=='completed')]
            else:
                scan_list = all_scans
            self.vprint("{info} Identified {new} scans to be processed".format(info=bcolors.INFO, new=len(scan_list)))

            if not scan_list:
                self.vprint("{info} No new scans to process. Exiting...".format(info=bcolors.INFO))
                exit(0)

            # Create scan subfolders
            for f in folders:
                if not os.path.exists(self.path_check(f['name'])):
                    if f['name'] == 'Trash' and self.nessus_trash:
                        os.makedirs(self.path_check(f['name']))
                    elif f['name'] != 'Trash':
                        os.makedirs(self.path_check(f['name']))
                else:
                    os.path.exists(self.path_check(f['name']))
                    self.vprint('{info} Directory already exist for {scan} - Skipping creation'.format(
                        scan=self.path_check(f['name']), info=bcolors.INFO))

            # try download and save scans into each folder the belong to
            scan_count = 0
            # TODO Rewrite this part to go through the scans that have aleady been processed
            for s in scan_list:
                scan_count += 1
                scan_name, scan_id, history_id,\
                norm_time, status, uuid = s['scan_name'], s['scan_id'], s['history_id'],\
                                          s['norm_time'], s['status'], s['uuid']

                # TODO Create directory sync function which scans the directory for files that exist already and populates the database

                folder_id = s['folder_id']
                scan_history = self.nessus.get_scan_history(scan_id)
                folder_name = next(f['name'] for f in folders if f['id'] == folder_id)
                if status == 'completed':
                    file_name = '%s_%s_%s_%s.%s' % (scan_name, scan_id, history_id, norm_time, 'csv')
                    repls = ('\\', '_'), ('/', '_'), ('/', '_'), (' ', '_')
                    file_name = reduce(lambda a, kv: a.replace(*kv), repls, file_name)
                    relative_path_name = self.path_check(folder_name + '/' + file_name)

                    if os.path.isfile(relative_path_name):
                        if self.develop:
                            csv_in = pd.read_csv(relative_path_name)
                            record_meta = (
                            scan_name, scan_id, norm_time, file_name, time.time(), csv_in.shape[0], 'nessus', uuid, 1)
                            self.record_insert(record_meta)
                            self.vprint(
                            "{info} File {filename} already exist! Updating database".format(info=bcolors.INFO, filename=relative_path_name))
                    else:
                        file_req = self.nessus.download_scan(scan_id=scan_id, history=history_id, export_format='csv')
                        clean_csv = pd.read_csv(io.StringIO(file_req.decode('utf-8')))
                        if len(clean_csv) > 2:
                            self.vprint("Processing %s/%s for scan: %s" % (scan_count, len(scan_history), scan_name))
                            clean_csv['CVSS'] = clean_csv['CVSS'].astype(str).apply(self.cleanser)
                            clean_csv['CVE'] = clean_csv['CVE'].astype(str).apply(self.cleanser)
                            clean_csv['Description'] = clean_csv['Description'].astype(str).apply(self.cleanser)
                            clean_csv['Synopsis'] = clean_csv['Description'].astype(str).apply(self.cleanser)
                            clean_csv['Solution'] = clean_csv['Solution'].astype(str).apply(self.cleanser)
                            clean_csv['See Also'] = clean_csv['See Also'].astype(str).apply(self.cleanser)
                            clean_csv['Plugin Output'] = clean_csv['Plugin Output'].astype(str).apply(self.cleanser)
                            clean_csv.to_csv(relative_path_name, index=False)
                            record_meta = (
                            scan_name, scan_id, norm_time, file_name, time.time(), clean_csv.shape[0], 'nessus', uuid,
                            1)
                            self.record_insert(record_meta)
                            self.vprint("{info} {filename} records written to {path} ".format(info=bcolors.INFO, filename=clean_csv.shape[0], path=file_name))
                        else:
                            record_meta = (
                            scan_name, scan_id, norm_time, file_name, time.time(), clean_csv.shape[0], 'nessus', uuid,
                            1)
                            self.record_insert(record_meta)
                            self.vprint(file_name + ' has no host available... Updating database and skipping!')
            self.conn.close()
            "{success} Scan aggregation complete! Connection to database closed.".format(success=bcolors.SUCCESS)


        else:
            self.vprint('{fail} Failed to use scanner at {host}'.format(fail=bcolors.FAIL, host=self.nessus_hostname+':'+self.nessus_port))