-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathkdi_rb_to_py.py
144 lines (110 loc) · 6 KB
/
kdi_rb_to_py.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
import json
import csv
from datetime import datetime
class KennaToolkit:
def __init__(self):
self.assets = []
self.vuln_defs = []
self.paged_assets = []
self.uploaded_files = []
def uniq(self, asset):
return {
"file": asset.get("file"),
"ip_address": asset.get("ip_address"),
"mac_address": asset.get("mac_address"),
"hostname": asset.get("hostname"),
"ec2": asset.get("ec2"),
"netbios": asset.get("netbios"),
"url": asset.get("url"),
"fqdn": asset.get("fqdn"),
"external_id": asset.get("external_id"),
"database": asset.get("database"),
"application": asset.get("application"),
"image": asset.get("image_id"),
"container": asset.get("container_id")
}
def kdi_initialize(self):
self.assets = []
self.vuln_defs = []
self.paged_assets = []
self.uploaded_files = []
def create_kdi_asset(self, asset_hash, dup_check=True):
self.kdi_initialize()
uniq_asset_hash = self.uniq(asset_hash)
if dup_check and any(uniq(a) == uniq_asset_hash for a in self.assets):
return None
asset_hash["tags"] = asset_hash.get("tags", [])
asset_hash["vulns"] = []
self.assets.append({k: v for k, v in asset_hash.items() if v is not None})
return {k: v for k, v in asset_hash.items() if v is not None}
def find_or_create_kdi_asset(self, asset_hash, match_key=None):
self.kdi_initialize()
uniq_asset_hash = self.uniq(asset_hash)
asset_hash_key = asset_hash.get(match_key) if match_key else None
a = next((asset for asset in self.assets if asset.get(match_key) == asset_hash_key), None) if match_key else None
if not a:
print("Unable to find asset {}, creating a new one...".format(asset_hash))
self.create_kdi_asset(asset_hash, False)
a = next((asset for asset in self.assets if asset.get(match_key) == asset_hash_key), None) if match_key else None
return a
def create_kdi_asset_vuln(self, asset_hash, vuln_hash, match_key=None):
self.kdi_initialize()
a = self.find_or_create_kdi_asset(asset_hash, match_key)
vuln_hash["status"] = vuln_hash.get("status", "open")
vuln_hash["port"] = int(vuln_hash.get("port", 0)) if vuln_hash.get("port") else None
now = datetime.utcnow().strftime("%Y-%m-%d")
vuln_hash["last_seen_at"] = vuln_hash.get("last_seen_at", now)
vuln_hash["created_at"] = vuln_hash.get("created_at", now)
a["vulns"] = a.get("vulns", [])
a["vulns"].append({k: v for k, v in vuln_hash.items() if v is not None})
return {k: v for k, v in vuln_hash.items() if v is not None}
def create_kdi_asset_finding(self, asset_hash, finding_hash, match_key=None):
self.kdi_initialize()
a = self.find_or_create_kdi_asset(asset_hash, match_key)
finding_hash["triage_state"] = finding_hash.get("triage_state", "new")
finding_hash["last_seen_at"] = finding_hash.get("last_seen_at", datetime.utcnow().strftime("%Y-%m-%d"))
a["findings"] = a.get("findings", [])
a["findings"].append({k: v for k, v in finding_hash.items() if v is not None})
return {k: v for k, v in finding_hash.items() if v is not None}
def create_paged_kdi_asset_vuln(self, asset_hash, vuln_hash, match_key=None):
self.kdi_initialize()
uniq_asset_hash = self.uniq(asset_hash)
asset_hash_key = asset_hash.get(match_key) if match_key else None
a = next((asset for asset in self.paged_assets if asset.get(match_key) == asset_hash_key), None) if match_key else None
if not a:
a = next((asset for asset in self.assets if asset.get(match_key) == asset_hash_key), None) if match_key else None
if a:
self.paged_assets.append(a)
self.assets.remove(a)
else:
a = asset_hash
self.paged_assets.append(a)
vuln_hash["status"] = vuln_hash.get("status", "open")
vuln_hash["port"] = int(vuln_hash.get("port", 0)) if vuln_hash.get("port") else None
vuln_hash["last_seen_at"] = vuln_hash.get("last_seen_at", datetime.utcnow().strftime("%Y-%m-%d"))
a["vulns"] = a.get("vulns", [])
a["vulns"].append({k: v for k, v in vuln_hash.items() if v is not None})
return True
def kdi_upload(self, output_dir, filename, kenna_connector_id, kenna_api_host, kenna_api_key, skip_autoclose=False, max_retries=3, version=1):
write_assets = self.paged_assets if self.paged_assets and any(self.paged_assets) else self.assets
if not write_assets:
return
self.write_file_stream(output_dir, filename, skip_autoclose, write_assets, self.vuln_defs, version)
print("Output is available at: {}/{}".format(output_dir, filename))
if kenna_connector_id and kenna_api_host and kenna_api_key:
print("Attempting to upload to Kenna API at {}".format(kenna_api_host))
response_json = self.upload_file_to_kenna_connector(kenna_connector_id, kenna_api_host, kenna_api_key, "{}/{}".format(output_dir, filename), False, max_retries)
filenum = response_json.get("data_file")
self.uploaded_files = [] if self.uploaded_files is None else self.uploaded_files
self.uploaded_files.append(filenum)
self.clear_data_arrays()
return response_json
def kdi_connector_kickoff(self, kenna_connector_id, kenna_api_host, kenna_api_key):
if not self.uploaded_files:
return
print("Attempting to run Kenna Connector at {}".format(kenna_api_host))
self.run_files_on_kenna_connector(kenna_connector_id, kenna_api_host, kenna_api_key, self.uploaded_files)
def clear_data_arrays(self):
self.assets = [] if self.paged_assets is None or not any(self.paged_assets) else self.paged_assets
self.paged_assets = []
self