#!/usr/bin/env python3

import logging
import json
from sys import argv, exit
import os
import scan_gvm_library as gvm_library
import scan_utilities as utilities
import argparse

### GVM Options ###
WAIT_TIMEOUT = 3600 #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.INFO,
    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 = utilities.import_dep_info(
        dep_json, 
        args.endpoint_keys)

logging.info(f"endpoints: {endpoints}")

# test gmp connection
gvm_library.set_auth(auth_name, auth_passwd)
logging.info(f"gvm version: {gvm_library.get_version()}")

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

    port_list = gvm_library.get_or_create_port_list(port_list_name,ports)
    logging.info(f"Port list:\n {utilities.pretty_json(port_list)}")
    
    target = gvm_library.get_or_create_target(target_name,host,port_list)
    logging.info(f"Target:\n {utilities.pretty_json(target)}")
    
    task = gvm_library.get_or_create_task(task_name, target)
    logging.info(f"Task:\n {utilities.pretty_json(task)}")
    
    if task['status'] == 'New':
        task = gvm_library.start_task(task)
    if gvm_library.wait_for_task_ending(task, WAIT_TIMEOUT):
        gvm_library.save_report(task,gvm_library.report_formats.txt, report_filename)
        reports[host] = gvm_library.get_report_info(task)
    else:
        reports[host] = f"ERROR Task: {task['id']}"
    
    gvm_library.delete_task(task)
    gvm_library.delete_target(target)
    gvm_library.delete_port_list(port_list)

reports = gvm_library.process_global_reports_info(reports)

logging.info(utilities.pretty_json(reports))

with open(summary_filename, "w") as f:
    f.write(json.dumps(reports))