from gvm.connections import TLSConnection
from gvm.protocols.gmpv208 import Gmp, AliveTest 
from gvm.transforms import EtreeTransform
from gvm.xml import pretty_print
from time import time, sleep
import logging
import json
import base64
from sys import argv, exit
import socket
import os

def get_version():
	with Gmp(connection, transform=transform) as gmp:
		gmp.authenticate(auth_name, auth_passwd)	
		pretty_print(gmp.get_version())

########## PORT LIST ##################################

def create_port_list(port_list_name, ports):
    with Gmp(connection, transform=transform) as gmp:
        gmp.authenticate(auth_name, auth_passwd)
        res = gmp.create_port_list(port_list_name, ','.join(ports))
        status = res.xpath('@status')[0]
        status_text = res.xpath('@status_text')[0]
        if status == "201":
            id = res.xpath('@id')[0]
            return {'name': port_list_name, 'id': id}
        else:
            msg = f"ERROR during Port list creation. Status code: {status}, msg: {status_text}"
            raise Exception(msg) 

def get_port_lists(filter_str="rows=-1"):
    l_o = []
    with Gmp(connection, transform=transform) as gmp:
        gmp.authenticate(auth_name, auth_passwd)
        res = gmp.get_port_lists(filter_string=filter_str)
        for pl in res.xpath('port_list'):
            o = dict()
            o['name'] = pl.xpath('name/text()')[0]
            o['id'] = pl.xpath('@id')[0]
            o['in_use'] = pl.xpath('in_use/text()')[0]
            l_o.append(o)
    return l_o

def delete_port_list(port_list):
    with Gmp(connection, transform=transform) as gmp:
        gmp.authenticate(auth_name, auth_passwd)
        res = gmp.delete_port_list(port_list['id'])
        status = res.xpath('@status')[0]
        status_text = res.xpath('@status_text')[0]
        if status == "200":
            print(f"Port_list with id: {port_list['id']} and name: {port_list['name']} DELETED") 
        else:
            print(f"ERROR {status}: {status_text}")

def get_or_create_port_list(port_list_name, ports):
    res = get_port_lists(port_list_name)
    if len(res) == 0:
        port_list = create_port_list(port_list_name, ports)
        return get_port_lists(port_list['id'])[0]
    elif len(res) == 1:
        return res[0]
    else:
        print(f"WARNING Found {len(res)} results.")
        return res

############## TARGET  ##################################

def create_target(name,ip,port_list,ovs_ssh_credential):
    o = dict()
    with Gmp(connection, transform=transform) as gmp:
        gmp.authenticate(auth_name, auth_passwd)
        res = gmp.create_target(
                name=name,
                comment = "",
                hosts=[ip],
                port_list_id = port_list['id'],
                ssh_credential_id = ovs_ssh_credential['id'],
                alive_test=AliveTest('Consider Alive'))
        status = res.xpath('@status')[0]
        status_text = res.xpath('@status_text')[0]
        if status == "201":
            id = res.xpath('@id')[0]
            return {'name': name, 'id': id}
        else:
            msg = f"ERROR during Target creation. Status code: {status}, msg: {status_text}"
            raise Exception(msg) 

def get_targets(filter_str):
    res = []
    with Gmp(connection, transform=transform) as gmp:
        gmp.authenticate(auth_name, auth_passwd)
        targets = gmp.get_targets(filter_string=filter_str)
        for target in targets.xpath('target'):
            o = dict()
            o['name'] = target.xpath('name/text()')[0]
            o['hosts'] = target.xpath('hosts/text()')[0]
            o['id'] = target.xpath('@id')[0]
            o['in_use'] = target.xpath('in_use/text()')[0]
            res.append(o)
    return res

def delete_target(target):
    with Gmp(connection, transform=transform) as gmp:
        gmp.authenticate(auth_name, auth_passwd)
        res = gmp.delete_target(target['id'])
        status = res.xpath('@status')[0]
        status_text = res.xpath('@status_text')[0]
        if status == "200":
            print(f"Port_list with id: {target['id']} and name: {target['name']} DELETED") 
        else:
            print(f"ERROR {status}: {status_text}")

def get_or_create_target(target_name,ip,port_list,ovs_ssh_credential):
    res = get_targets(target_name)
    if len(res) == 0:
        t = create_target(target_name,ip,port_list,ovs_ssh_credential)
        return get_targets(t['id'])[0]
    elif len(res) == 1:
        return res[0]
    else:
        print(f"Found {len(res)} results. Return None")
        return res

def search_and_delete_target(target_name):
    targets = get_targets(target_name)
    if len(targets) == 1:
        delete_target(targets[0]['id'])
    else:
        raise("Multiple results for search")

def search_and_delete_all_targets(target_name):
    targets = get_targets(target_name)
    for target in targets:
        delete_target(target)

############## TASK ##################################

def create_task(name, config, target, scanner):
    o = dict()
    with Gmp(connection, transform=transform) as gmp:
        gmp.authenticate(auth_name, auth_passwd)
        res = gmp.create_task(
                name=name,
                config_id=config['id'],
                target_id=target['id'],
                scanner_id=scanner['id'])
        status = res.xpath('@status')[0]
        status_text = res.xpath('@status_text')[0]
        if status == "201":
            id = res.xpath('@id')[0]
            return {'name': name, 'id': id}
        else:
            msg = f"ERROR during Task creation. Status code: {status}, msg: {status_text}"
            raise Exception(msg)

def get_tasks(filter_str):
    res = []
    with Gmp(connection, transform=transform) as gmp:
        gmp.authenticate(auth_name, auth_passwd)
        tasks = gmp.get_tasks(filter_string=filter_str)
        for task in tasks.xpath('task'):
                o = dict()
                o['name'] = task.xpath('name/text()')[0]
                o['id'] = task.xpath('@id')[0]
                o['progress'] = task.xpath('progress/text()')[0]
                o['in_use'] = task.xpath('in_use/text()')[0]
                o['status'] = task.xpath('status/text()')[0]
                o['target_id'] = task.xpath('target/@id')[0]
                try:
                    o['report_id'] = task.xpath('last_report/report/@id')[0]
                except:
                    pass
                res.append(o)
    return res

def get_or_create_task(task_name, config, target, scanner):
    res = get_tasks(task_name)
    if len(res) == 0:
        t = create_task(task_name, config, target, scanner)
        return get_tasks(t['id'])[0]
    elif len(res) == 1:
        return res[0]
    else:
        print(f"Found {len(res)} results. Return None")
        return res

def get_all_tasks():
    res = []
    with Gmp(connection, transform=transform) as gmp:
        gmp.authenticate(auth_name, auth_passwd)
        tasks = gmp.get_tasks(filter_string="rows=-1")
        for task in tasks.xpath('task'):
                o = dict()
                o['name'] = task.xpath('name/text()')[0]
                o['id'] = task.xpath('@id')[0]
                o['progress'] = task.xpath('progress/text()')[0]
                o['in_use'] = task.xpath('in_use/text()')[0]
                o['status'] = task.xpath('status/text()')[0]
                o['target_id'] = task.xpath('target/@id')[0]
                try:
                    o['report_id'] = task.xpath('last_report/report/@id')[0]
                except:
                    pass
                res.append(o)
    return res

def search_and_delete_all_tasks(filter_str):
    tasks = get_tasks(filter_str)
    for task in tasks:
        delete_task(task)

def start_task(task):
    with Gmp(connection, transform=transform) as gmp:
        gmp.authenticate(auth_name, auth_passwd)
        res = gmp.start_task(task['id'])
        task['report_id'] = res.xpath('report_id/text()')[0]
    return task        

def stop_task(task):
    with Gmp(connection, transform=transform) as gmp:
        gmp.authenticate(auth_name, auth_passwd)
        pretty_print(gmp.stop_task(task['id']))

def delete_task(task):
    with Gmp(connection, transform=transform) as gmp:
        gmp.authenticate(auth_name, auth_passwd)
        res = gmp.delete_task(task['id'])
        status = res.xpath('@status')[0]
        status_text = res.xpath('@status_text')[0]
        if status == "200":
            print(f"Target with id: {task['id']} and name: {task['name']} DELETED") 
        else:
            print(f"ERROR {status}: {status_text}")

############## REPORTS #####################################3

class report_formats:
    anonymous_xml = "5057e5cc-b825-11e4-9d0e-28d24461215b"
    csv_results   = "c1645568-627a-11e3-a660-406186ea4fc5"
    itg           = "77bd6c4a-1f62-11e1-abf0-406186ea4fc5"
    pdf           = "c402cc3e-b531-11e1-9163-406186ea4fc5"
    txt           = "a3810a62-1f62-11e1-9219-406186ea4fc5"
    xml           = "a994b278-1f62-11e1-96ac-406186ea4fc5"

def get_report_formats():
    with Gmp(connection, transform=transform) as gmp:
        gmp.authenticate(auth_name, auth_passwd)   
        res =  gmp.get_report_formats()
        for f in res.xpath('report_format'):
            name = f.xpath('name/text()')[0]
            id = f.xpath('@id')[0]
            print(id,name)

def get_report_format(id):
    with Gmp(connection, transform=transform) as gmp:
        gmp.authenticate(auth_name, auth_passwd)   
        res =  gmp.get_report_formats()
        pretty_print(res)  

def get_progress(task):
    task_info = get_tasks(task['id'])[0]
    status = task_info['status']         #   New -> Requested -> Queued -> Running  -> Done
    progress = int(task_info['progress'])#    0         0           0      0 -> 100     -1
    return status, progress

def wait_for_task_completition(task, timeout=3600):
    start_time = time()
    while True:
        status, progress = get_progress(task)
        if status not in ["New","Requested","Queued","Running","Done"]: # ["Interrupted", ...]
            return False
        if status == "Done" and progress == -1:
            return True
        if time() - start_time > timeout:
                print("TIMEOUT during waiting for task completition")
                return False
        sleep(10)

def save_report(task,report_format_id, report_filename ):
    with Gmp(connection, transform=transform) as gmp:
        gmp.authenticate(auth_name, auth_passwd)
        res = gmp.get_report(task['report_id'],
                                    report_format_id=report_format_id, 
                                    ignore_pagination=True,
                                    details="1")
        code = str(res.xpath('report/text()')[0])
        with open(report_filename, "wb") as fh:
            fh.write(base64.b64decode(code))

def save_severity_report(task, severity_filename):
    dict_severity = {"Log": 0, "Low": 1, "Medium": 2, "High": 3}
    with Gmp(connection, transform=transform) as gmp:
        gmp.authenticate(auth_name, auth_passwd)
        res = gmp.get_report(task['report_id'],
                            report_format_id=report_formats.anonymous_xml, 
                            ignore_pagination=True,
                            details="1")
        severities = res.xpath('report/report/ports/port/threat/text()')
        old_num_severity = 0
        severity = "Log"
        for sev in severities:
            if dict_severity[sev] > old_num_severity:
                old_num_severity = dict_severity[sev]
                severity = sev
        with open(severity_filename, "w") as f:
            f.write(severity)

def get_severity(task):
    dict_severity = {"Log": 0, "Low": 1, "Medium": 2, "High": 3}
    with Gmp(connection, transform=transform) as gmp:
        gmp.authenticate(auth_name, auth_passwd)
        res = gmp.get_report(task['report_id'],
                            report_format_id=report_formats.anonymous_xml, 
                            ignore_pagination=True,
                            details="1")
        severities = res.xpath('report/report/ports/port/threat/text()')
        old_num_severity = 0
        severity = "Log"
        for sev in severities:
            if dict_severity[sev] > old_num_severity:
                old_num_severity = dict_severity[sev]
                severity = sev
        return severity
            
def get_reports(filter_str="rows=-1"):
    lo = []
    with Gmp(connection, transform=transform) as gmp:
        gmp.authenticate(auth_name, auth_passwd)
        reports = gmp.get_reports(filter_string = filter_str)
        for report in reports.xpath('report'):
            o = dict()
            o['task_name'] = report.xpath('task/name/text()')[0]
            o['id'] = report.xpath('@id')[0]
            lo.append(o)
        return lo



def get_numeric_severity(severity):
    if severity == "Log":
        return 0
    elif severity == "Low":
        return 1
    elif severity == "Medium":
        return 2
    elif severity == "High":
        return 3
    else:
        return 4

def get_severity_from_number(num):
    if num == 0:
        return "Low"
    elif num == 1:
        return "Low"
    elif num == 2:
        return "Medium"
    elif num == 3:
        return "High"
    else:
        return "Undefined"

def process_global_severity(severities):
    max_num_severity = 0
    for _,sev in severities.items():
        numeric_severity = get_numeric_severity(sev)
        if numeric_severity > max_num_severity:
            max_num_severity = numeric_severity
    severities['global_severity'] = get_severity_from_number(max_num_severity)
    if max_num_severity < 2:
        severities['global'] = "OK"
    else:
        severities['global'] = "NOK"
    return severities

def print_pretty_json(j):
    print(json.dumps(j,sort_keys=True,indent=4))
        
def import_dep_info(file_path, endpoints_to_scan):
    with open(file_path) as f:
        data = json.load(f)

    endpoints = dict()
    for key in data['outputs'].keys():
        if key in endpoints_to_scan:
            endpoint = str(data['outputs'][key]).split("/")[2]
            print("endpoint",endpoint)
            host,port = endpoint.split(':')
            if host not in endpoints:
                endpoints[host] = {"22"}
            endpoints[host].add(port)
    return endpoints
                
################ MAIN #######################################

auth_name = os.getenv('GMP_USER')
auth_passwd = os.getenv('GMP_PASSWORD')
print(auth_name, auth_passwd)
logging.basicConfig(filename='debug.log', level=logging.DEBUG)
#local_ip = socket.gethostbyname(socket.gethostname())
local_ip = '127.0.0.1'
connection = TLSConnection(hostname=local_ip)
transform = EtreeTransform()
config = {'id':"9866edc1-8869-4e80-acac-d15d5647b4d9"}
scanner = {'id': "08b69003-5fc2-4037-a479-93b440211c73"}
ovs_ssh_credential = {'id': "a89d5ebf-01bf-4836-ae72-a65bc633219a"}
wait_timeout = 3600 #1h

if len(argv) != 4:
    print("Please pass three parameters:")
    print("- endpoints to scans [endpoints1,endpoint2,endpoint3,...]")
    print("- dep.json path [/home/gmp/workspace/dep.json]")
    print("- output directory [/home/gmp/workspace]")
    exit(1)

endpoints_to_scan = argv[1].split(',')
dep_json = argv[2]
output_dir = argv[3]
print("endpoints_to_scan", endpoints_to_scan)
print("dep_json", dep_json)
print("output_dir", output_dir)

endpoints = import_dep_info(dep_json, endpoints_to_scan)
#print_pretty_json(endpoints)      
severities = dict()
for host,ports in endpoints.items():
    print(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"
    severity_filename = f"{output_dir}/severity.json"

    port_list = get_or_create_port_list(port_list_name,ports)
    print_pretty_json(port_list)
    target = get_or_create_target(target_name,host,port_list,ovs_ssh_credential)
    print_pretty_json(target)
    task = get_or_create_task(task_name, config, target,scanner)
    print_pretty_json(task)
    if task['status'] == 'New':
        task = start_task(task)
    if wait_for_task_completition(task,wait_timeout):
        save_report(task,report_formats.txt, report_filename)
        severities[host] = get_severity(task)
    else:
        severities[host] = f"ERROR Task: {task['id']}"
    delete_task(task)
    delete_target(target)
    delete_port_list(port_list)

severities = process_global_severity(severities)
with open(severity_filename, "w") as f:
    f.write(json.dumps(severities))