#!/usr/bin/env python3

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

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

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(
    "--output-dir", 
    default="."
    )

args = parser.parse_args()

logging.basicConfig(
    filename='scan.log', 
    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")

dep_json = args.dep_json
template_url = args.template_url
output_dir = args.output_dir
dep_name = (template_url.split('/')[-1]).split('.')[0]
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}")

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 issues irrelevant for INFN
issues_to_drop = read_not_relevant_issues()

tasks = list()
targets = list()
port_lists = list()

reports = 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"
    summary_filename = f"{output_dir}/summary-report.json"

    # Create PortList obj related to endpoint
    port_list = PortList(client = gvm.get_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.get_client(),
                    name = target_name, 
                    host = host, 
                    port_list =port_list)
    logging.info(f"Target:\n {target}")
    
    task = Task(client= gvm.get_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")
        reports[host] = task.get_report_info(issues_to_drop)
    else:
        reports[host] = {'global': {"severity": -1, "threat": f"Scan Error. task.id: {task.id}"} }

reports = process_global_reports_info(reports)
logging.info(pretty_json(reports))

# Write global reports on file
with open(summary_filename, "w") as f:
    f.write(json.dumps(reports))  

for task in tasks:
    task.delete()

for target in targets:
    target.delete()
    
for port_list in port_lists:
    port_list.delete()