#!/usr/bin/env python3

import logging
import json
import os
from gvm_library import GVMClient, ReportFormats, pretty_json
from utilities import import_dep_info, process_global_reports_info, read_not_relevant_issues
import argparse

### GVM Options ###
WAIT_TIMEOUT = 7200 #1h

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(
    "--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 Exception("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 Exception("GMP_PASSWORD env var is not defined")

dep_json = args.dep_json
output_dir = args.output_dir
logging.info(f"endpoint_keys: {args.endpoint_keys}")
logging.info(f"dep_json: {dep_json}")
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_client = GVMClient(auth_n = auth_name, auth_p = auth_passwd)

# Print GVM library version
logging.info(f"gvm version: {gvm_client.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_{host}"
    task_name = f"{auth_name}_task_{host}"
    port_list_name = f"{auth_name}_pl_{host}"
    report_filename = f"{output_dir}/{host}-report"
    summary_filename = f"{output_dir}/summary-report.json"

    port_list = gvm_client.get_or_create_port_list(port_list_name, ports)
    logging.info(f"Port list:\n {port_list}")
    
    target = gvm_client.get_or_create_target(target_name, host, port_list)
    logging.info(f"Target:\n {target}")
    
    task = gvm_client.get_or_create_task(task_name, 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}"} }

for task in tasks:
    try:
        del task 
    except Exception as e:
        logging.error(f"Error during task deletion. Msg: {e.msg}")

for target in targets:
    try:
        del target 
    except Exception as e:
        logging.error(f"Error during target deletion. Msg: {e.msg}")

for port_list in port_lists:
    try:
        del port_list 
    except Exception as e:
        logging.error(f"Error during port list deletion. Msg: {e.msg}")

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))