-
Notifications
You must be signed in to change notification settings - Fork 1
/
autorun.py
196 lines (166 loc) · 6.49 KB
/
autorun.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
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
#!/usr/bin/python3
import argparse
import threading
import queue
from datetime import datetime
import time
import xml.etree.ElementTree as ET
import mysql.connector
import os
import settings
import subprocess
exit_flag = 0
error_plugin_list = list()
class my_thread (threading.Thread):
def __init__(self, thread_name, working_queue):
threading.Thread.__init__(self)
self.thread_name = thread_name
self.working_queue = working_queue
def run(self):
process_data(self.thread_name, self.working_queue)
def process_data(thread_name, working_queue):
while not exit_flag:
queue_lock.acquire()
if not work_queue.empty():
work = working_queue.get()
plugin_id = work['plugin_id']
target_host = work['target_host']
target_port = work['target_port']
target_protocol = work['target_protocol']
message = '%s for Plugin ID: %s; Host: %s; Port: %s\n%s' % (thread_name, plugin_id, target_host, target_port, time_elapsed())
print(good_msg(message))
# Start processing the stuff here
conn = settings.db()
mycursor = conn.cursor(prepared=True)
sql = 'SELECT * FROM `plugin` INNER JOIN `command` ON `plugin`.`command_id` = `command`.`command_id` WHERE `plugin`.`plugin_id` = %s'
mycursor.execute(sql, (plugin_id, ))
result = mycursor.fetchall()
parent_directory = autorun_output_file_path + "/" + str(plugin_id) + "/"
if not (os.path.isdir(parent_directory)):
os.mkdir(parent_directory)
if mycursor.rowcount == 0:
if plugin_id not in error_plugin_list:
error_plugin_list.append(plugin_id)
error(plugin_id, target_host, target_port, target_protocol)
error_file = parent_directory + "ERROR_" + str(plugin_id) + ".txt"
if not (os.path.isfile(error_file)):
f = open(error_file, "a+")
f.write("Error Plugin ID: " + str(plugin_id) + " is missing from database.")
f.close()
else:
if result[0][1] != 1:
if isinstance(result[0][3], str):
command = result[0][3] + ' 2>&1'
else:
command = result[0][3].decode('utf-8') + ' 2>&1'
destination_text_file_path = parent_directory + str(plugin_id)+'_' + str(target_host)+'_'+str(target_port) + '.txt'
if '\r' in command:
command = command.replace('\r', '')
if '\n' in command:
command = command.replace('\n', '')
if 'IP_ADDRESS' in command:
command = command.replace('IP_ADDRESS', str(target_host))
if 'PORT' in command:
command = command.replace('PORT', str(target_port))
if 'PLUGIN' in command:
command = command.replace('PLUGIN', str(plugin_id))
if 'RANDOM_PATH' in command:
target_txt_file = '%s_%s_%s_TEMP.txt' % (plugin_id, target_host, target_port)
random_path = '"' + parent_directory + target_txt_file + '"'
command = command.replace('RANDOM_PATH', random_path)
task_left = work_queue.qsize()
task_completed = total_tasks - task_left
if not (os.path.exists(destination_text_file_path)):
# Do not change this part, create a file for better debugging when a thread hangs/wont terminate
f = open(destination_text_file_path, "w")
f.write("[Command] " + command + '\n\n')
f.write(subprocess.check_output(command, shell=True).decode())
f.close()
message = 'Completed command for Plugin ID: %s; Host: %s; Port: %s; Task Left: %s; Task Completed: %s\n%s' % (plugin_id, target_host, target_port, task_left, task_completed, time_elapsed())
print(good_msg(message))
else:
message = 'Skipped executing Plugin ID: %s because the validation check already exists for specific host: %s:%s\n%s' % (plugin_id, target_host, target_port, time_elapsed())
print(warning_msg(message))
mycursor.close()
conn.close()
# End processing the stuff here
queue_lock.release()
time.sleep(1)
def error(plugin_id, target_host, target_port, target_protocol):
f = open('plugin_check_output.txt', 'a+')
f.close()
if str(plugin_id) not in open('plugin_check_output.txt').read():
f = open('plugin_check_output.txt', 'a+')
row = str(plugin_id)+' - '+str(target_host)+':'+str(target_port)+' ('+str(target_protocol).upper()+')\r\n'
f.write(row)
f.close()
message = 'Error Detected on Plugin ID '+str(plugin_id)+', please check on plugin_check_output.txt for more information\n'+time_elapsed()
print(error_msg(message))
def good_msg(message):
message = '\033[92m'+message+'\033[0m'
return message
def warning_msg(message):
message = '\033[93m'+message+'\033[0m'
return message
def error_msg(message):
message = '\033[91m'+message+'\033[0m'
return message
def time_elapsed():
end_time = datetime.now()
time_taken = end_time - start_time
message = 'Total Time Elapsed: '+str(time_taken)
return message
parser = argparse.ArgumentParser()
parser.add_argument('-f', metavar='Nessus File', type=str, help='raw Nessus file to do verification', required=True)
parser.add_argument('-t', metavar='Threads', type=int, help='no of threads to run', required=False, default=10)
args = parser.parse_args()
nessus_file = args.f
total_thread = args.t
total_tasks = 0
start_time = datetime.now()
message = "Autorun started on: " + str(start_time)
print(good_msg(message))
output_directory_name = os.path.splitext(os.path.basename(nessus_file))[0]
directory_count = 1
while True:
autorun_output_file_path = output_directory_name + "_" + str(directory_count)
directory_count += 1
if not (os.path.isdir(autorun_output_file_path)):
break
os.mkdir(autorun_output_file_path)
queue_lock = threading.Lock()
work_queue = queue.Queue(maxsize=0)
threads = []
# Fill the queue
queue_lock.acquire()
tree = ET.parse(nessus_file)
root = tree.getroot()[1]
total_hosts = len(root)
for i in range(0, total_hosts):
target_host = root[i].attrib['name']
total_findings = len(root[i])
for j in range(1, total_findings):
findings = list()
risk_factor = root[i][j].find('risk_factor').text
if (risk_factor != 'None'):
plugin_id = root[i][j].attrib['pluginID']
target_port = root[i][j].attrib['port']
target_protocol = root[i][j].attrib['protocol']
total_tasks += 1
work_queue.put({'plugin_id': plugin_id, 'target_host': target_host, 'target_port': target_port, 'target_protocol': target_protocol})
queue_lock.release()
# Create new threads
for i in range(0, total_thread):
thread_name = 'Thread-'+str(i)
thread = my_thread(thread_name, work_queue)
thread.start()
threads.append(thread)
# Wait for queue to empty
while not work_queue.empty():
pass
# Notify threads it's time to exit
exit_flag = 1
# Wait for all threads to complete
for t in threads:
t.join()
print(good_msg('Completed!\n' + time_elapsed()))