class MainWindow(Frame):
def __init__(self, master=None):
super().__init__(master)
self.db_manager = DBManager()
self.vt_api = VirusTotalApi()
self.file_manager = FileManager()
self.popup_is_alive = False
self.master = master
self.master.protocol("WM_DELETE_WINDOW", self.on_closing)
self.master.minsize(width=800, height=600)
self.master.title("PVBPS")
self.create_widgets()
self.scan_thread = threading.Thread()
def create_widgets(self):
#Grid.
self.master.grid_columnconfigure(0, weight=0)
self.master.grid_columnconfigure(1, weight=0)
self.master.grid_columnconfigure(2, weight=1)
self.master.grid_columnconfigure(3, weight=0)
self.master.grid_rowconfigure(0, weight=0)
self.master.grid_rowconfigure(1, weight=1)
self.master.grid_rowconfigure(2, weight=0)
self.master.resizable(True, True)
#Toolbar.
self.menu = Menu(self.master)
self.master.config(menu=self.menu)
self.file = Menu(self.menu)
self.file.add_command(label="Scan folder...", command=self.init_scan)
self.file.add_command(label="Exit", command=self.on_closing)
self.menu.add_cascade(label="File", menu=self.file)
#Description labels.
self.description_frame = LabelFrame(self.master, text="Description")
self.description_frame.grid(row=0,
column=0,
columnspan=4,
sticky=(N, E, W))
self.lbl_md5 = Label(self.description_frame, text="MD5:")
self.lbl_md5.grid(row=0,
column=0,
padx=(5, 5),
pady=(5, 5),
sticky=(W))
#self.lbl_md5.config(height = 5)
self.lbl_ratio = Label(self.description_frame, text="Detection ratio:")
self.lbl_ratio.grid(row=1,
column=0,
padx=(5, 5),
pady=(5, 5),
sticky=(W))
#self.lbl_ratio.config(height = 5)
self.lbl_date = Label(self.description_frame, text="Analysis date:")
self.lbl_date.grid(row=2,
column=0,
padx=(5, 5),
pady=(5, 5),
sticky=(W))
#self.lbl_date.config(height = 5)
self.lbl_md5_val = Label(self.description_frame, text="")
self.lbl_md5_val.grid(row=0,
column=1,
padx=(5, 5),
pady=(5, 5),
sticky=(W))
self.lbl_ratio_val = Label(self.description_frame, text="")
self.lbl_ratio_val.grid(row=1,
column=1,
padx=(5, 5),
pady=(5, 5),
sticky=(W))
self.lbl_date_val = Label(self.description_frame, text="")
self.lbl_date_val.grid(row=2,
column=1,
padx=(5, 5),
pady=(5, 5),
sticky=(W))
#Scan progressbar.
self.lbl_scan = Label(self.description_frame, text="Scan progress:")
self.lbl_scan.grid(row=3,
column=0,
padx=(5, 5),
pady=(5, 5),
sticky=(W))
self.progressbar = progressbar = Progressbar(self.description_frame,
orient=HORIZONTAL,
length=180,
mode='indeterminate')
self.progressbar.grid(row=4,
column=0,
padx=(5, 5),
pady=(5, 5),
sticky=(W, E))
#File listbox.
self.file_list = Listbox(self.master)
self.file_list.config(width=30)
self.file_list.bind('<<ListboxSelect>>', self.on_file_select)
#Scrollbars for listbox.
self.file_list_xscrollbar = Scrollbar(self.master, orient=HORIZONTAL)
self.file_list_yscrollbar = Scrollbar(self.master, orient=VERTICAL)
self.file_list.configure(xscrollcommand=self.file_list_xscrollbar.set)
self.file_list.configure(yscrollcommand=self.file_list_yscrollbar.set)
self.file_list_xscrollbar.configure(command=self.file_list.xview)
self.file_list_yscrollbar.configure(command=self.file_list.yview)
self.file_list.grid(row=1, column=0, sticky=(W, N, S))
self.file_list_xscrollbar.grid(row=2,
column=0,
columnspan=2,
sticky=(E, W))
self.file_list_yscrollbar.grid(row=1,
rowspan=2,
column=1,
sticky=(N, S))
#Treeview.
self.tree = Treeview(self.master)
self.tree.grid(row=1, column=2, sticky=(N, S, E, W))
self.tree.heading('#0', text='Antivirus')
self.tree["columns"] = ("one", "two")
self.tree.column("one", width=200)
self.tree.column("two", width=200)
self.tree.heading("one", text="Result")
self.tree.heading("two", text="Update")
#Treeview scrollbars.
self.tree_xscrollbar = Scrollbar(self.master, orient=HORIZONTAL)
self.tree_yscrollbar = Scrollbar(self.master, orient=VERTICAL)
self.tree.configure(xscrollcommand=self.tree_xscrollbar.set)
self.tree.configure(yscrollcommand=self.tree_yscrollbar.set)
self.tree_xscrollbar.configure(command=self.tree.xview)
self.tree_yscrollbar.configure(command=self.tree.yview)
self.tree_xscrollbar.grid(row=2, column=2, columnspan=1, sticky=(E, W))
self.tree_yscrollbar.grid(row=1, rowspan=1, column=3, sticky=(N, S))
def on_closing(self):
if messagebox.askokcancel("Quit", "Do you want to exit?"):
#I think you don't wanna wait...
#so I comment this out.
#if self.scan_thread.is_alive():
# messagebox.showinfo("Scan is in progress!", "Please wait untill the scan is finished!")
#else:
self.master.destroy()
def on_file_select(self, evt):
w = evt.widget
index = int(w.curselection()[0])
value = w.get(index)
self.update_tree(self.file_manager.files[value])
self.update_labels(self.file_manager.files[value])
def update_file_list(self):
#Delete previous scan.
self.file_list.delete(0, END)
for key, value in self.file_manager.files.items():
self.file_list.insert(END, key)
def set_no_result(self):
self.lbl_ratio_val["text"] = "None"
self.lbl_date_val["text"] = "None"
self.lbl_md5_val["text"] = "None"
def update_labels(self, selected_record):
record = self.db_manager.get_record(selected_record)
if not record:
self.set_no_result()
else:
record = record[0]
self.lbl_ratio_val["text"] = "{}/{}".format(
record["positives"], record["total"])
self.lbl_date_val["text"] = record["scan_date"]
self.lbl_md5_val["text"] = record["md5"]
def update_tree(self, selected_record):
#Delete previous.
for i in self.tree.get_children():
self.tree.delete(i)
#Show selected.
key_scans = "scans"
row_index = 0
record = self.db_manager.get_record(selected_record)
if not record:
self.set_no_result()
else:
for key, value in record[0]["scans"].items():
self.tree.insert("",
row_index,
text=key,
values=(value["result"], value["update"]))
row_index = row_index + 1
def ask_to_upload(self, not_analyzed):
not_analyzed_count = len(not_analyzed)
if not_analyzed_count > 0:
if messagebox.askokcancel(
"Some files weren't analyzed!",
"Some files weren't found in the database, do you want to upload them (Warning: this may take a while!), for later scan?"
):
#Upload.
self.vt_api.upload_files(not_analyzed)
def ask_to_encrypt(self, malicious):
malicious_count = len(malicious)
if malicious_count > 0:
if messagebox.askokcancel(
"Malware found!",
"Do you want to encrypt these malicious files?"):
#Encypt.
self.file_manager.encrypt_files(malicious)
else:
messagebox.showinfo("No malware found!",
"No Malware was found in scanned files.")
def check_scan_progress(self):
if not self.scan_thread.is_alive():
self.progressbar.stop()
messagebox.showinfo("Done!", "Scan is complete!")
else:
self.after(100, self.check_scan_progress)
def init_scan(self):
#Delete cached records.
self.file_manager.reset_files()
self.db_manager.reset_records()
try:
self.file_manager.directory = filedialog.askdirectory()
if self.file_manager.is_dir_ok():
self.scan_thread = threading.Thread(target=self.scan_folder,
args=())
self.scan_thread.start()
self.progressbar.start()
messagebox.showinfo(
"Started...",
"Scan has started, see progressbar on the left side. Please wait for the scan to finish!"
)
self.check_scan_progress()
except IOError:
print("Error! Unable to open this directory. Please try again...")
def finish_scan(self):
#Ask to encrypt malicious files.
self.ask_to_encrypt(self.db_manager.get_malicious())
#Ask to upload files, that werent sucessfuly analyzed via hash.
self.ask_to_upload(
self.db_manager.get_not_analyzed(self.file_manager.files))
def scan_folder(self):
self.file_manager.directory_files()
self.file_manager.compute_md5_hash()
#Show in GUI.
self.update_file_list()
#Check local DB.
#Store MD5 to check with the API.
to_check = []
#new_records = []
for f, md5 in self.file_manager.files.items():
record = self.db_manager.get_record(md5)
#If record is not in the local DB, check it later with with the API.
if record == None:
to_check.append(md5)
else:
#Cache record.
self.db_manager.cache_record(record)
self.db_manager.get_malicious()
#If there are some MD5s to check with the API, do it.
if len(to_check) > 0:
self.vt_api.get_batch_response(to_check,
self.db_manager.add_records)
#Add new records to the DB.
#And cache them.
#self.db_manager.add_records(new_records)
#Finish scan.
self.finish_scan()
