Skip to content

Commit 4a9708b

Browse files
diStyAppsdiStyApps
committed
0.1.5
some refactoring and fixes for future updates.
1 parent cb7390f commit 4a9708b

File tree

10 files changed

+386
-401
lines changed

10 files changed

+386
-401
lines changed

‎.gitignore‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,3 +133,6 @@ dmypy.json
133133

134134
# Pyre type checker
135135
.pyre/
136+
137+
#vscode
138+
settings.json

‎CONSTANTS.py‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
2+
SCANNED_FILES = 'SCANNED FILES: '
3+
INFECTED_FILES = 'INFECTED FILES: '
4+
DANGEROUS_GLOBALS = 'DANGEROUS GLOBALS: '
5+
SCANNED_FILES_DEF = f'{SCANNED_FILES} 0'
6+
INFECTED_FILES_DEF = f'{INFECTED_FILES} 0'
7+
DANGEROUS_GLOBALS_DEF = f'{DANGEROUS_GLOBALS} 0'
8+
9+
SCANNED_FILES_TXT_KEY = '-scanned_files_txt-'
10+
INFECTED_FILES_TXT_KEY = '-infected_files_txt-'
11+
DANGEROUS_GLOBALS_TXT_KEY = '-dangerous_globals_txt-'
12+
13+
HUGGINGFACE_RAD_KEY = '-huggingface_radio-'
14+
URL_RAD_KEY = '-url_radio-'
15+
FILE_RAD_KEY = '-file_radio-'
16+
DIRECTORY_RAD_KEY = '-directory_radio-'
17+
TYPE_SELECTOR_INP_RAD_KEY = '-type_selector_input_radio-'
18+
19+
FILE_FRM_KEY = '-file_frame-'
20+
DIRECTORY_FRM_KEY = '-directory_frame-'
21+
URL_FRM_KEY = '-url_frame-'
22+
HUGGINGFACE_FRM_KEY = '-huggingface_frame-'
23+
24+
HUGGINGFACE_INP_KEY = '-huggingface_inp_files-'
25+
URL_INP_KEY = '-url_inp_files-'
26+
FILE_INP_KEY = '-file_inp_files-'
27+
DIRECTORY_INP_KEY = '-directory_inp_files-'
28+
29+
SCAN_BTN_KEY = '-scan_btn-'
30+
SCAN_BTN_LBL = 'SCAN'
31+
32+
HUGGINGFACE_LBL = 'Hugging Face'
33+
URL_LBL = 'URL'
34+
FILE_LBL = 'File'
35+
DIRECTORY_LBL = 'Directory'
36+
37+
CONSOLE_ML_KEY = '-console_ml-'
38+
STATUS_INFO_TXT_KEY = '-status_info_txt-'
39+
STATUS_STATE_TXT_KEY = '-status_state_txt-'
40+
41+
HUGGINGFACE_CLEAR_BTN_KEY = '-huggingface_clear_btn-'
42+
URL_CLEAR_BTN_KEY = '-url_clear_btn-'
43+
44+
FONT = 'Arial 12'
45+

‎picklescan_gui.py‎

Lines changed: 0 additions & 400 deletions
This file was deleted.

‎run_app_gui.py‎

Lines changed: 252 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,252 @@
1+
2+
import PySimpleGUI as sg
3+
from picklescan.src.picklescan.scanner import scan_directory_path, scan_file_path,scan_huggingface_model, scan_url, scanned_files, infected_files
4+
from util.ui_flattener import flatten_ui_elements
5+
from util.clipboard import do_clipboard_operation, right_click_menu
6+
from util.file_extension import file_ext
7+
import util.support as support
8+
import util.colors as color
9+
from CONSTANTS import *
10+
11+
sg.theme('Dark Gray 15')
12+
__version__ = '0.1.5'
13+
APP_TITLE = f"Disty's Stable Diffusion Pickle Scanner GUI - Ver {__version__}"
14+
isError:int = 0
15+
16+
def main():
17+
#region layout
18+
19+
def browse_layout(key_,visible_,disabled=False):
20+
if key_ == 'huggingface':
21+
browse_type = sg.Button('Clear',k=f'-{key_}_clear_btn-',disabled=disabled,size=(10,2))
22+
if key_ == 'url':
23+
browse_type = sg.Button('Clear',k=f'-{key_}_clear_btn-',disabled=disabled,size=(10,2))
24+
if key_ == 'file':
25+
browse_type = sg.FileBrowse(k=f'-{key_}_FileBrowse-',file_types=(file_ext),disabled=disabled,size=(10,2))
26+
if key_ == 'directory':
27+
browse_type = sg.FolderBrowse(k=f'-{key_}_FolderBrowse-',disabled=disabled,size=(10,2))
28+
29+
layout = sg.Frame('',[
30+
[
31+
sg.Input(key=f'-{key_}_inp_files-',enable_events=True,expand_x=True,expand_y=True,font=FONT,
32+
background_color=color.DARK_GRAY,right_click_menu=right_click_menu),
33+
browse_type,
34+
],
35+
],expand_x=True,k=f'-{key_}_frame-',visible=visible_,relief=sg.RELIEF_SOLID,border_width=1,background_color=color.GRAY_9900)
36+
return layout
37+
38+
top_column = [
39+
[
40+
sg.Frame('',[
41+
[
42+
support.buttons_layout(),
43+
],
44+
[
45+
sg.Radio(HUGGINGFACE_LBL,TYPE_SELECTOR_INP_RAD_KEY,default=True,k=HUGGINGFACE_RAD_KEY,enable_events=True),
46+
sg.Radio(URL_LBL,TYPE_SELECTOR_INP_RAD_KEY,default=False,k=URL_RAD_KEY,enable_events=True),
47+
sg.Radio(FILE_LBL,TYPE_SELECTOR_INP_RAD_KEY,default=False,k=FILE_RAD_KEY,enable_events=True),
48+
sg.Radio(DIRECTORY_LBL,TYPE_SELECTOR_INP_RAD_KEY,default=False,k=DIRECTORY_RAD_KEY,enable_events=True),
49+
],
50+
],expand_x=True,expand_y=False,relief=sg.RELIEF_SOLID,border_width=1,visible=True,background_color=color.GRAY_9900)
51+
],
52+
[
53+
browse_layout('huggingface',True),
54+
browse_layout('url',False),
55+
browse_layout('file',False),
56+
browse_layout('directory',False),
57+
],
58+
]
59+
60+
mid_column = [
61+
[
62+
sg.Frame('',[
63+
[
64+
sg.Button(SCAN_BTN_LBL,k=SCAN_BTN_KEY,font=FONT,expand_x=True,size=(30,2),mouseover_colors=(color.GRAY_9900,color.DARK_GREEN)),
65+
],
66+
[
67+
sg.Text(SCANNED_FILES_DEF,k=SCANNED_FILES_TXT_KEY,expand_x=True,expand_y=True,font=FONT,justification='left',size=(20,1)),
68+
sg.Text(INFECTED_FILES_DEF,k=INFECTED_FILES_TXT_KEY,expand_x=True,expand_y=True,font=FONT,justification='c',size=(20,1)),
69+
sg.Text(DANGEROUS_GLOBALS_DEF,k=DANGEROUS_GLOBALS_TXT_KEY,expand_x=True,expand_y=True,font=FONT,justification='right',size=(20,1)),
70+
],
71+
[
72+
sg.Multiline(k=CONSOLE_ML_KEY,reroute_stdout=True,write_only=False,reroute_cprint=True,
73+
autoscroll=True,border_width=0,sbar_width=20,sbar_trough_color=0,
74+
background_color=color.GRAY_1111, auto_refresh=True,size=(120,30),expand_x=True,expand_y=True)
75+
],
76+
[
77+
sg.Frame('',[
78+
[
79+
sg.Text('',key=STATUS_STATE_TXT_KEY, expand_x=True,expand_y=True,justification='c',font=FONT,background_color=color.GRAY_9900)
80+
],
81+
],expand_x=True,expand_y=False,relief=sg.RELIEF_FLAT,border_width=0,visible=True,element_justification='c',background_color=color.GRAY_9900,size=(50,30))
82+
],
83+
],expand_x=True,expand_y=True,border_width=0,relief=sg.RELIEF_FLAT,element_justification="c")
84+
],
85+
]
86+
87+
bottom_column = [
88+
[
89+
sg.Frame('',[
90+
[
91+
sg.Frame(' Status',[
92+
[
93+
sg.Text('',key=STATUS_INFO_TXT_KEY, expand_x=True,font=FONT)
94+
],
95+
],expand_x=True,expand_y=True,relief=sg.RELIEF_SOLID,border_width=0,visible=True,element_justification='c',background_color=color.GRAY_9900,title_color=color.DARK_BLUE)
96+
],
97+
98+
],expand_x=True,border_width=0,relief=sg.RELIEF_FLAT,element_justification='c')
99+
],
100+
]
101+
102+
layout = [
103+
top_column ,
104+
mid_column,
105+
bottom_column,
106+
]
107+
108+
#endregion layout
109+
110+
window = sg.Window(APP_TITLE,layout,finalize=True, resizable=True,enable_close_attempted_event=False,background_color=color.GRAY_9900)
111+
112+
#region window elements
113+
114+
huggingface_inp_elem:sg.Input = window[HUGGINGFACE_INP_KEY]
115+
url_inp_elem:sg.Input = window[URL_INP_KEY]
116+
file_inp_elem:sg.Input = window[FILE_INP_KEY]
117+
directory_inp_elem:sg.Input = window[DIRECTORY_INP_KEY]
118+
status_state_txt_elem:sg.Text = window[STATUS_STATE_TXT_KEY]
119+
scanned_files_txt_elem:sg.Text = window[SCANNED_FILES_TXT_KEY]
120+
infected_files_txt_elem:sg.Text = window[INFECTED_FILES_TXT_KEY]
121+
dangerous_globals_txt_elem:sg.Text = window[DANGEROUS_GLOBALS_TXT_KEY]
122+
console_ml_elem:sg.Multiline = window[CONSOLE_ML_KEY]
123+
status_info_txt_elem:sg.Text = window[STATUS_INFO_TXT_KEY]
124+
file_frame_elem:sg.Frame = window[FILE_FRM_KEY]
125+
directory_frame_elem:sg.Frame = window[DIRECTORY_FRM_KEY]
126+
url_frame_elem:sg.Frame = window[URL_FRM_KEY]
127+
huggingface_frame_elem:sg.Frame = window[HUGGINGFACE_FRM_KEY]
128+
129+
#endregion window elements
130+
131+
selected_input = {
132+
HUGGINGFACE_RAD_KEY: {file_frame_elem: False, directory_frame_elem: False, url_frame_elem: False, huggingface_frame_elem: True},
133+
URL_RAD_KEY: {file_frame_elem: False, directory_frame_elem: False, url_frame_elem: True, huggingface_frame_elem: False},
134+
FILE_RAD_KEY: {file_frame_elem: True, directory_frame_elem: False, url_frame_elem: False, huggingface_frame_elem: False},
135+
DIRECTORY_RAD_KEY: {file_frame_elem: False, directory_frame_elem: True, url_frame_elem: False, huggingface_frame_elem: False},
136+
}
137+
138+
flatten_ui_elements(window)
139+
140+
def reset_ui():
141+
reset_files_state_display()
142+
update_status_state('')
143+
console_ml_elem.update(value='')
144+
scanned_files.clear()
145+
infected_files.clear()
146+
147+
def update_status_state(value):
148+
status_state_txt_elem.update(value=value)
149+
status_state_txt_elem.Widget.config(background=color.GRAY_9900)
150+
status_state_txt_elem.ParentRowFrame.config(background=color.GRAY_9900)
151+
152+
def reset_files_state_display():
153+
status_info_txt_elem.update(value='Scanning. Please wait...')
154+
scanned_files_txt_elem.update(value=SCANNED_FILES_DEF)
155+
infected_files_txt_elem.update(value=INFECTED_FILES_DEF)
156+
dangerous_globals_txt_elem.update(value=DANGEROUS_GLOBALS_DEF)
157+
158+
def scan():
159+
input_files = None
160+
if values[HUGGINGFACE_RAD_KEY]:
161+
input_files = values[HUGGINGFACE_INP_KEY]
162+
scan_target = scan_huggingface_model
163+
elif values[URL_RAD_KEY]:
164+
input_files = values[URL_INP_KEY]
165+
scan_target = scan_url
166+
elif values[DIRECTORY_RAD_KEY]:
167+
input_files = values[DIRECTORY_INP_KEY]
168+
scan_target = scan_directory_path
169+
elif values[FILE_RAD_KEY]:
170+
input_files = values[FILE_INP_KEY]
171+
scan_target = scan_file_path
172+
173+
if input_files is None or len(input_files) <= 1:
174+
status_info_txt_elem.update(value='Nothing to scan')
175+
return
176+
177+
try:
178+
return scan_target(input_files)
179+
except (RuntimeError, ConnectionRefusedError, AttributeError, FileNotFoundError,UnboundLocalError) as e:
180+
print(e.__class__.__name__, e)
181+
status_info_txt_elem.update(value='Fail to scan')
182+
183+
def update_results_display(scan_result):
184+
scanned_files_txt_elem.update(value=f'{SCANNED_FILES} {scan_result.scanned_files}')
185+
infected_files_txt_elem.update(value=f'{INFECTED_FILES} {scan_result.infected_files}')
186+
dangerous_globals_txt_elem.update(value=f'{DANGEROUS_GLOBALS} {scan_result.issues_count}')
187+
scanned_files_len = len(scanned_files)
188+
infected_files_len = len(infected_files)
189+
190+
print('')
191+
if scanned_files_len > 0:
192+
if infected_files_len == 0:
193+
status_state_txt_elem.update(text_color='white')
194+
status_state_txt_elem.Widget.config(background=color.DARK_GREEN)
195+
status_state_txt_elem.ParentRowFrame.config(background=color.DARK_GREEN)
196+
status_state_txt_elem.update(value='PASS')
197+
198+
if infected_files_len > 0:
199+
print('')
200+
print('----------- INFECTED FILES -----------')
201+
print('')
202+
203+
for file in infected_files:
204+
print(f" {file}")
205+
206+
status_state_txt_elem.update(text_color='white')
207+
status_state_txt_elem.Widget.config(background=color.RED_ORANGE)
208+
status_state_txt_elem.ParentRowFrame.config(background=color.RED_ORANGE)
209+
status_state_txt_elem.update(value='FAIL')
210+
211+
status_info_txt_elem.update(value='Done')
212+
213+
while True:
214+
event, values = window.read()
215+
216+
if event == sg.WIN_CLOSED:
217+
break
218+
219+
if event == SCAN_BTN_KEY:
220+
reset_ui()
221+
scan_result = scan()
222+
223+
if not isError:
224+
if scan_result:
225+
update_results_display(scan_result)
226+
227+
if event == HUGGINGFACE_CLEAR_BTN_KEY:
228+
huggingface_inp_elem.update(value='')
229+
230+
if event == URL_CLEAR_BTN_KEY:
231+
url_inp_elem.update(value='')
232+
233+
if event in selected_input:
234+
for frame, visibility in selected_input[event].items():
235+
frame.update(visible=visibility)
236+
237+
if event in right_click_menu[1]:
238+
if values[HUGGINGFACE_RAD_KEY]:
239+
do_clipboard_operation(event, window, huggingface_inp_elem,STATUS_INFO_TXT_KEY)
240+
if values[URL_RAD_KEY]:
241+
do_clipboard_operation(event, window, url_inp_elem,STATUS_INFO_TXT_KEY)
242+
if values[FILE_RAD_KEY]:
243+
do_clipboard_operation(event, window, file_inp_elem,STATUS_INFO_TXT_KEY)
244+
if values[DIRECTORY_RAD_KEY]:
245+
do_clipboard_operation(event, window, directory_inp_elem,STATUS_INFO_TXT_KEY)
246+
247+
window[CONSOLE_ML_KEY].update(value='')
248+
249+
support.buttons(event)
250+
251+
if __name__ == '__main__':
252+
main()

‎util/clipboard.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
import PySimpleGUI as sg
2+
right_click_menu = ['', ['Copy', 'Paste', 'Cut', 'Clear']]
3+
4+
def do_clipboard_operation(event, window, element,STATUS_INFO_TXT_KEY):
5+
if event == 'Select All':
6+
element.Widget.selection_clear()
7+
# element.Widget.tag_add('sel', '1.0', 'end')
8+
elif event == 'Copy':
9+
try:
10+
text = element.Widget.selection_get()
11+
window.TKroot.clipboard_clear()
12+
window.TKroot.clipboard_append(text)
13+
except:
14+
# print('Nothing selected')
15+
window[STATUS_INFO_TXT_KEY].update(value='Nothing selected')
16+
17+
elif event == 'Paste':
18+
element.Widget.insert(sg.tk.INSERT, window.TKroot.clipboard_get())
19+
elif event == 'Cut':
20+
try:
21+
text = element.Widget.selection_get()
22+
window.TKroot.clipboard_clear()
23+
window.TKroot.clipboard_append(text)
24+
element.update('')
25+
except:
26+
# print('Nothing selected')
27+
window[STATUS_INFO_TXT_KEY].update(value='Nothing selected')
28+
elif event == 'Clear':
29+
try:
30+
text = element.Widget.selection_get()
31+
element.update('')
32+
except:
33+
# print('Nothing selected')
34+
window[STATUS_INFO_TXT_KEY].update(value='Nothing selected')

‎util/colors.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
DARK_GREEN = '#78BA04'
2+
GRAY_9900 = '#0A0A0A'
3+
DARK_GRAY = '#1F1F1F'
4+
GRAY_1111 = '#111111'
5+
TERMINAL_BLUE = '#0099ff'
6+
TERMINAL_GREEN = '#00ff00'
7+
DARK_BLUE = '#4974a5'
8+
RED_ORANGE = '#C13515'

‎util/file_extension.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
file_ext = {
2+
("All", "*.ckpt"),
3+
("All", "*.pth"),
4+
("All", "*.pt"),
5+
("All", "*.bin"),
6+
("All", "*.pkl"),
7+
("All", "*.pickle"),
8+
("All", "*.joblib"),
9+
("All", "*.dat"),
10+
("All", "*.data"),
11+
("All", "*.npy"),
12+
("All", "*.npy"),
13+
("All", "*.npz"),
14+
("All", "*.zip"),
15+
}

0 commit comments

Comments
 (0)