#!/usr/bin/env python3

import logging
import json
import os
from gvm_library import GVMClient, ReportFormats, GvmException
from gvm_library import PortList, Task, Target
from gvm_library import ReportManager
from utilities import import_dep_info
import argparse

### GVM Options ###
WAIT_TIMEOUT = 7200 #2h

## TEMPLATE CONSTANTs ###
OS_DISTR = "os_distribution"
OS_VERS = "os_version"

## FILENAMES
SUMMARY_FILENAME = "summary-report.json"
OIDS_FILENAME = "oids.yaml"
LOG_FILENAME = "log_scan.log"



parser = argparse.ArgumentParser(
    description='Scan endpoints and machines')
            
parser.add_argument(
    "--endpoint-keys",
    help="Orchestrator output endpoints to scan (endpoints1,endpoints2)",
    default="None"
    )
parser.add_argument(
    "--dep-json", 
    default="./dep.json"
    )

parser.add_argument(
    "--template-url", 
    default="None"
    )

parser.add_argument(
    "--template-input",
    default="None"
    )

parser.add_argument(
    "--output-dir", 
    default="."
    )

args = parser.parse_args()

logging.basicConfig(
    filename = LOG_FILENAME,
    level = logging.DEBUG,
    format = '%(asctime)s %(levelname)-8s %(message)s',
    datefmt = '%Y-%m-%d %H:%M:%S',
    filemode = 'w')
logging.info("\n\nStart scan application")

if os.environ.get('GMP_USER') is not None and \
        os.environ.get('GMP_USER') != '':
    auth_name = os.getenv('GMP_USER')
else:
    logging.error("GMP_USER env var is not defined\nexit")
    raise GvmException("GMP_USER env var is not defined")

if os.environ.get('GMP_PASSWORD') is not None and \
        os.environ.get('GMP_PASSWORD') != '':
    auth_passwd = os.getenv('GMP_PASSWORD')
else:
    logging.error("GMP_PASSWORD env var is not defined\nexit")
    raise GvmException("GMP_PASSWORD env var is not defined")

# Import passed parameters
dep_json = args.dep_json
template_url = args.template_url
template_input = args.template_input
output_dir = args.output_dir

dep_name = (template_url.split('/')[-1]).split('.')[0]
summary_filename = os.path.join(output_dir, SUMMARY_FILENAME)
oids_filename = os.path.join(output_dir, OIDS_FILENAME)

# Log imported parameters
logging.info(f"endpoint_keys: {args.endpoint_keys}")
logging.info(f"dep_json: {dep_json}")
logging.info(f"template_url: {template_url}")
logging.info(f"dep_name: {dep_name}")
logging.info(f"output_dir: {output_dir}")

d_input = json.loads(template_input)
os_name = None
if OS_DISTR in d_input and OS_VERS in d_input:
    os_name = f"{d_input[OS_DISTR]}{d_input[OS_VERS]}"
logging.info(f"os_name: {os_name}")

# Check if this script will scans an operating system
#is_os: bool = "single_vm" in dep_name and os_name is not None
is_os: bool = os_name is not None
logging.info(f"is_os: {is_os}")

endpoints = import_dep_info(dep_json, args.endpoint_keys)
logging.info(f"endpoints: {endpoints}")

# Create GVM client
gvm = GVMClient(auth_n = auth_name, auth_p = auth_passwd)

# Print GVM library version
logging.info(f"gvm version: {gvm.get_version()}")

# Retrieve os related issues
report_manager = ReportManager(os_name, is_os)

# init variables
tasks = list()
targets = list()
port_lists = list()
reports = dict()
oids = dict()

for host,ports in endpoints.items():
    logging.info(f"endpoint: {host}:{ports}")
    
    target_name = f"{auth_name}_target_{dep_name}_{host}"
    task_name = f"{auth_name}_task_{dep_name}_{host}"
    port_list_name = f"{auth_name}_pl_{dep_name}_{host}"
    report_filename = f"{output_dir}/{host}-report"

    # Create PortList obj related to endpoint
    port_list = PortList(client = gvm.client, 
                         name = port_list_name, 
                         ports = ports)
    logging.info(f"Port list:\n {port_list}")
    
    # Create Target obj related to endpoint
    target = Target(client = gvm.client,
                    name = target_name, 
                    host = host, 
                    port_list = port_list)
    logging.info(f"Target:\n {target}")
    
    task = Task(client = gvm.client,
                name = task_name,
                target = target)
    logging.info(f"Task:\n {task}")

    tasks.append(task)
    targets.append(target)
    port_lists.append(port_list)
    
    if task.status == 'New':
        task.start()
    if task.wait(WAIT_TIMEOUT):
        task.save_report(format = ReportFormats.pdf, 
                         filename = f"{report_filename}.pdf")
        task.save_report(format = ReportFormats.txt, 
                         filename = f"{report_filename}.txt")
        report_manager.import_report(host,task.get_report())
    else:
        task_report = {"severity": ReportManager.MAX_SEVERITY, 
                       "threat": f"Scan Error. task.id: {task.id}"}
        report_manager.import_report(host,task_report)

    #task.delete()
    #target.delete()
    #port_list.delete()

report_manager.classify_reports()
logging.info("Reports")
logging.info(report_manager.get_summary())
logging.info("Classified oids")
logging.info(report_manager.get_classified_oids())

# Write reports on file
report_manager.write_data(summary_filename,
                          oids_filename)