Minor Bug fixes

utils/script.py

@@ -2,17 +2,16 @@
 #auth_name = 'backbone'
 #auth_passwd = 'krowfodwas3olrab'
 
-from hashlib import new
 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
-from datetime import datetime
 import json
 import base64
 from sys import argv, exit
+import socket
 
 def get_version():
 	with Gmp(connection, transform=transform) as gmp:
@@ -71,7 +70,7 @@ def get_or_create_port_list(port_list_name, ports):
 
 ############## TARGET  ##################################
 
-def create_target(name,ip,port_list,ovs_ssh_credential_id):
+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)
@@ -80,7 +79,7 @@ def create_target(name,ip,port_list,ovs_ssh_credential_id):
                 comment = "",
                 hosts=[ip],
                 port_list_id = port_list['id'],
-                ssh_credential_id = ovs_ssh_credential_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]
@@ -116,10 +115,10 @@ def delete_target(target):
         else:
             print(f"ERROR {status}: {status_text}")
 
-def get_or_create_target(target_name,ip,port_list,ovs_ssh_credential_id):
+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_id)
+        t = create_target(target_name,ip,port_list,ovs_ssh_credential)
         return get_targets(t['id'])[0]
     elif len(res) == 1:
         return res[0]
@@ -141,15 +140,15 @@ def search_and_delete_all_targets(target_name):
 
 ############## TASK ##################################
 
-def create_task(name, config_id, target_id, scanner_id):
+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)
+                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":
@@ -179,10 +178,10 @@ def get_tasks(filter_str):
                 res.append(o)
     return res
 
-def get_or_create_task(task_name, config_id, target_id, scanner_id):
+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_id, target_id, scanner_id)
+        t = create_task(task_name, config, target, scanner)
         return get_tasks(t['id'])[0]
     elif len(res) == 1:
         return res[0]
@@ -240,6 +239,14 @@ def delete_task(task):
 
 ############## 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():
     # |------------- ID -----------------| |--- NAME ---|
     # 5057e5cc-b825-11e4-9d0e-28d24461215b Anonymous XML
@@ -297,7 +304,7 @@ def save_severity_report(task, severity_filename):
     with Gmp(connection, transform=transform) as gmp:
         gmp.authenticate(auth_name, auth_passwd)
         res = gmp.get_report(task['report_id'],
-                            report_format_id="5057e5cc-b825-11e4-9d0e-28d24461215b", 
+                            report_format_id=report_formats.anonymous_xml, 
                             ignore_pagination=True,
                             details="1")
         severities = res.xpath('report/report/ports/port/threat/text()')
@@ -315,7 +322,7 @@ def get_severity(task):
     with Gmp(connection, transform=transform) as gmp:
         gmp.authenticate(auth_name, auth_passwd)
         res = gmp.get_report(task['report_id'],
-                            report_format_id="5057e5cc-b825-11e4-9d0e-28d24461215b", 
+                            report_format_id=report_formats.anonymous_xml, 
                             ignore_pagination=True,
                             details="1")
         severities = res.xpath('report/report/ports/port/threat/text()')
@@ -339,8 +346,7 @@ def get_reports(filter_str="rows=-1"):
             lo.append(o)
         return lo
 
-def print_pretty_json(j):
-    print(json.dumps(j,sort_keys=True,indent=4))
+
 
 def get_numeric_severity(severity):
     if severity == "Log":
@@ -378,50 +384,55 @@ def process_global_severity(severities):
     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 = "jenkins"
 auth_passwd = "bleyrauvHecsUbDy"
 logging.basicConfig(filename='debug.log', level=logging.DEBUG)
-connection = TLSConnection(hostname='172.17.0.2')
+local_ip = socket.gethostbyname(socket.gethostname())
+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"
-txt_report_format_id = "a3810a62-1f62-11e1-9219-406186ea4fc5"
-csv_report_format_id = "c1645568-627a-11e3-a660-406186ea4fc5"
-xml_report_format_id = "5057e5cc-b825-11e4-9d0e-28d24461215b"
+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
-print("len_argv:",len(argv))
-for i in range(len(argv)):
-    print(i,argv[i])
-  
-endpoints = argv[1].split(',')
+
+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", endpoints)
+print("endpoints_to_scan", endpoints_to_scan)
 print("dep_json", dep_json)
 print("output_dir", output_dir)
 
-with open(dep_json) as f:
-    data = json.load(f)
-
-endpoints = dict()
-for key in data['outputs'].keys():
-    if "endpoint" in key:
-        endpoint = str(data['outputs'][key]).split("://")[1]
-        print("endpoint",endpoint)
-        host,port = endpoint.split(':')
-        if host not in endpoints:
-            endpoints[host] = ["22"]
-        endpoints[host].append(port)
-
-print(json.dumps(endpoints,sort_keys=True,indent=4))      
+endpoints = import_dep_info(dep_json, endpoints_to_scan)
+print_pretty_json(endpoints)      
 severities = dict()
-
 for host,ports in endpoints.items():
     print(host,ports)
   
@@ -431,16 +442,16 @@ for host,ports in endpoints.items():
     report_filename = f"{output_dir}/{host}-report.txt"
     severity_filename = f"{output_dir}/severity.json"
 
-    pl = get_or_create_port_list(port_list_name,ports)
-    print_pretty_json(pl)
-    target = get_or_create_target(target_name,host,pl,ovs_ssh_credential_id)
+    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_id, target['id'],scanner_id)
+    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,txt_report_format_id, report_filename)
+        save_report(task,report_formats.txt, report_filename)
         severities[host] = get_severity(task)
     else:
         severities[host] = f"ERROR Task: {task['id']}"