#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
########################################
#                                      #
#  Used to control Nessus through      #
#  scripting.                          #
#                                      #
########################################

from nessus.mod_nessus.funciones_nessus import *
import requests
import json
import time
import sys
from termcolor import colored
import getpass
import os


def get_last_history(scan_id):
    data = connect('GET', '/scans/{0}'.format(scan_id))
    return data['history'][get_last_history_index(data['history'])]['history_id']


def get_last_history_index(history):
    count = 0
    for hist in history:
        count = count + 1
    return count - 1


def export_path(scan_id):
        data = {"format": "pdf", "chapters": "vuln_hosts_summary", "history_id": "{0}".format(get_last_history(scan_id))}
        data = connect('POST', '/scans/{0}/export'.format(scan_id), data=data)
        file_id = data['file']
        while export_status(scan_id, file_id) is False:
                time.sleep(5)
        return file_id


# Funcion parecida a la otra funcion download de funciones_nessus
def download_path(report_name, scan_id, file_id, save_path):
    if not (os.path.exists(save_path)):
        os.mkdir(save_path)

    data = connect('GET', '/scans/{0}/export/{1}/download'.format(scan_id, file_id))
    file_name = '{0}_{1}.pdf'.format(report_name, file_id)

    print('Saving report to {0}/{1}'.format(save_path, file_name))
    with open(os.path.join(save_path, file_name.replace('\t', '')), 'wb') as f:
        f.write(data)


def count_scan(scans, folder_id):
    count = 0
    for scan in scans:
        if scan['folder_id'] == folder_id:
            count = count + 1
    return count


def get_folder_id(serch_folder_name, data):
        folder_id = 0
        for folder in data['folders']:
                if folder['name'] == serch_folder_name:
                        folder_id = folder['id']
                        break
        return folder_id


def export_folder(data, folder_name = 'all'):
        if folder_name == 'All' or folder_name == 'all':
                for scan in data['scans']:
                        file_id = export_path(scan['id'])
                        download_path(scan['name'], scan['id'], file_id, os.path.join(os.getcwd(), folder_name))
        else:
                folder_id = get_folder_id(folder_name, data)
                if count_scan(data['scans'], folder_id) == 0:
                        print("Reports to download is missing")
                        return
                if folder_id != 0:
                        for scan in data['scans']:
                                if scan['folder_id'] == folder_id:
                                        file_id = export_path(scan['id'])
                                        download_path(scan['name'], scan['id'], file_id, os.path.join(os.getcwd(), folder_name))
                else:
                        print("{0} folder not found".format(folder_name))


def print_scans(data, folder_name='all'):
    if folder_name == 'All' or folder_name == 'all':
        for folder in data['folders']:
            print("\\{0} - ({1})\\".format(folder['name'], count_scan(data['scans'], folder['id'])))
            for scan in data['scans']:
                if scan['folder_id'] == folder['id']:
                    print("\t\"{0}\" - uuid: {1}".format(scan['name'], scan['uuid']))
    else:
        folder_id = get_folder_id(folder_name, data)
        if count_scan(data['scans'], folder_id) == 0:
            print("\\{0} - Folder is empty\\".format(folder_name))
            return
        if folder_id != 0:
            print("\\{0} - ({1})\\".format(folder_name, count_scan(data['scans'], folder_id)))
            for scan in data['scans']:
                if scan['folder_id'] == folder_id:
                    print("\t\"{0}\"".format(scan['name']))
        else:
            print("{0} folder not found".format(folder_name))


def list_scan():
        data = connect('GET', '/scans')
        return data


def listado_informes(token):
    # Mostramos los informes disponibles
    print(colored("\nList of reports...", 'cyan'))
    rep_list = list_scan()
    print_scans(rep_list)

    # Exportamos los directorios de Nessus que nos interesen
    print(colored("Exporting reports...", 'green'))
    exp_folder_name = input(colored('Input folder name to export (type "all" to export all reports): ', 'yellow'))
    export_folder(rep_list, exp_folder_name)


def exporta_escaneo_pdf():
    while True:
        global token
        print(colored("\nExportar PDF", 'cyan'))
        # Obtenemos el token para acceder al sistema Nessus
        token = obtener_token_nessus()

        # Mostramos los informes disponibles llamando al funcion listado_informes
        listado_informes(token)

        repite_consulta = input(colored("\nDesea repetir la consulta? (s/n): ", 'magenta', attrs=['bold']))
        if repite_consulta == 'n':
            break


def listado_directivas():
    # Funcion para mostrar el contenido de la directivas disponibles en nexus
    global token
    print(colored("\nListando directivas disponibles.", 'cyan'))
    print("")

    # Obtenemos el token para acceder al sistema Nessus
    token = obtener_token_nessus()
    # Llamamos a la funcion pasandole el token de conexion
    listar_directivas_existentes()

    input(colored("\nPulse INTRO para continua.", 'red'))
