fix #3 skips file scan and notfy to console about Corrupted file

diStyApps · diStyApps · commit ca58440f9662 · 2022-11-25T19:29:17.000+02:00
fix scan button not visible moved up
add some redsign
diff --git a/GUI.py b/GUI.py
@@ -9,6 +9,7 @@
 from picklescan.src.picklescan.scanner import scanned_files
 from picklescan.src.picklescan.scanner import infected_files
 import util.icons as ic
+from zipfile import BadZipFile
 COLOR_DARK_GREEN = '#78BA04'
 COLOR_DARK_BLUE = '#4974a5'
 COLOR_RED_ORANGE = '#C13515'
@@ -48,25 +49,40 @@ def do_clipboard_operation(event, window, element):
         except:
             # print('Nothing selected')
             window['-status_info-'].update(value='Nothing selected') 
+
 def main():
-    ver = '0.1.0'
+    ver = '0.1.1'
     sg.theme('Dark Gray 15')
     app_title = f"Disty's Stable Diffusion Pickle Scanner GUI - Ver {ver}"
     isError = 0
     file_ext = {
-        ("Pytorch Files", "*.ckpt"),
-        ("Pytorch Files", "*.pth"),
-        ("Pytorch Files", "*.pt"),
-        ("Pytorch Files", "*.bin"),
-        ("Pickle Files", "*.pkl"),
-        ("Pickle Files", "*.pickle"),
-        ("Pickle Files", "*.joblib"),
-        ("Pickle Files", "*.dat"),
-        ("Pickle Files", "*.data"),
-        ("Pickle Files", "*.npy"),
-        ("Pickle Files", "*.npy"),
-        ("Zip Files", "*.npz"),
-        ("Zip Files", "*.zip"),
+
+        ("All", "*.ckpt"),
+        ("All", "*.pth"),
+        ("All", "*.pt"),
+        ("All", "*.bin"),
+        ("All", "*.pkl"),
+        ("All", "*.pickle"),
+        ("All", "*.joblib"),
+        ("All", "*.dat"),
+        ("All", "*.data"),
+        ("All", "*.npy"),
+        ("All", "*.npy"),
+        ("All", "*.npz"),
+        ("All", "*.zip"),
+        # ("Pytorch Files", "*.ckpt"),
+        # ("Pytorch Files", "*.pth"),
+        # ("Pytorch Files", "*.pt"),
+        # ("Pytorch Files", "*.bin"),
+        # ("Pickle Files", "*.pkl"),
+        # ("Pickle Files", "*.pickle"),
+        # ("Pickle Files", "*.joblib"),
+        # ("Pickle Files", "*.dat"),
+        # ("Pickle Files", "*.data"),
+        # ("Pickle Files", "*.npy"),
+        # ("Pickle Files", "*.npy"),
+        # ("Zip Files", "*.npz"),
+        # ("Zip Files", "*.zip"),        
     }
 
     #region layout
@@ -83,7 +99,8 @@ def browse_layout(type_,visible_,disabled=False):
         layout = sg.Frame('',[
                                 [
                                     sg.Input(key=f'-input_files_{type_}-',enable_events=True,expand_x=True,expand_y=True,font='Ariel 11',background_color=COLOR_DARK_GRAY,right_click_menu=right_click_menu),
-                                    browse_type
+                                    browse_type,
+
                                 ],
                             ],expand_x=True,k=f'-frame_{type_}-',visible=visible_,relief=sg.RELIEF_SOLID,border_width=1,background_color=COLOR_GRAY_9900)
         return layout
@@ -118,44 +135,52 @@ def browse_layout(type_,visible_,disabled=False):
         ],
     ]
 
-    bottom_column = [
-        [
-            sg.Frame('',[
-                    [
-                        sg.Frame(' Status',[
-                            [
-                                sg.Text('',key='-status_info-', expand_x=True)
-                            ],
-                        ],expand_x=True,expand_y=False,relief=sg.RELIEF_SOLID,border_width=0,visible=True,element_justification='c',background_color=COLOR_GRAY_9900,title_color=COLOR_DARK_BLUE)                          
-                    ],  
-                [              
-                    sg.Button('Scan',k='-scan_button-',disabled=False,expand_x=True),
-                ],
-            ],expand_x=True,border_width=0,relief=sg.RELIEF_FLAT,element_justification='c')
-        ],    
-    ]
 
     console_column = [
         [
             sg.Frame('',[
+                                [              
+                        sg.Button('SCAN',k='-scan_button-',disabled=False,button_color=("white",COLOR_DARK_BLUE),font='Ariel 12 ',expand_x=True,size=(10,2)),
+                    ], 
+                    
                 [
                     sg.Text(SCANNED_FILES_DEF,k='-scanned_files_text-',expand_x=True,expand_y=True,font='Ariel 12',justification='left',size=(20,1)),
                     sg.Text(INFECTED_FILES_DEF,k='-infected_files_text-',expand_x=True,expand_y=True,font='Ariel 12',justification='c',size=(20,1)),
                     sg.Text(DANGEROUS_GLOBALS_DEF,k='-dangerous_globals_text-',expand_x=True,expand_y=True,font='Ariel 12',justification='right',size=(20,1)),
-                ],         
+                ], 
+  
+                    [
+                        sg.MLine(k='-console_ml-',reroute_stdout=True,write_only=False,reroute_cprint=True, autoscroll=True, text_color='white', auto_refresh=True,size=(120,30),expand_x=True,expand_y=True)
+                    ],                                              
                     [
                         sg.Frame('',[
                             [
                                 sg.Text('',key='-status_state_text-', expand_x=True,expand_y=True,justification='c',font='Ariel 12 bold',background_color=COLOR_GRAY_9900)
                             ],
                         ],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))                          
-                    ],                           
-                [
-                    sg.MLine(k='-console_ml-',reroute_stdout=True,write_only=False,reroute_cprint=True, autoscroll=True, text_color='white', auto_refresh=True,size=(120,35),expand_x=True,expand_y=True)
-                ],
-            ],expand_x=True,expand_y=True,border_width=0,relief=sg.RELIEF_FLAT,k='-scanned_frame-')
+                    ],   
+                                                           
+
+            ],expand_x=True,expand_y=True,border_width=0,relief=sg.RELIEF_FLAT,k='-scanned_frame-',element_justification="c")
         ],  
     ]
+    
+    bottom_column = [
+        [
+            sg.Frame('',[
+                    [
+                        sg.Frame(' Status',[
+                            [
+                                sg.Text('',key='-status_info-', expand_x=True,font='Ariel 12 bold')
+                            ],
+                        ],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)                          
+                    ],  
+                # [              
+                #     sg.Button('Scan',k='-scan_button-',disabled=False,expand_x=True),
+                # ],
+            ],expand_x=True,border_width=0,relief=sg.RELIEF_FLAT,element_justification='c')
+        ],    
+    ]
 
     layout = [
 
@@ -274,7 +299,6 @@ def browse_layout(type_,visible_,disabled=False):
                             print(ConnectionRefusedError,e)   
                             isError = 1
                             status_info_widget.update(value='Fail to scan') 
-
                         except AttributeError as e:
                             print(AttributeError,e)   
                             isError = 1  
@@ -292,7 +316,6 @@ def browse_layout(type_,visible_,disabled=False):
                             print(FileNotFoundError,e)   
                             isError = 1
                             status_info_widget.update(value='Fail to scan') 
-
                     else:
                             isError = 1
                             status_info_widget.update(value='Nothing to scan') 
@@ -302,45 +325,54 @@ def browse_layout(type_,visible_,disabled=False):
                     if len(input_files)>1:
                         try:
                             scan_result = scan_file_path(input_files)
+                            if scan_result == None:
+                                isError = 1
+                                status_info_widget.update(value='Fail to scan')                                
                         except FileNotFoundError as e:
                             print(FileNotFoundError,e)   
                             isError = 1
-                            status_info_widget.update(value='Fail to scan') 
+                            status_info_widget.update(value='Fail to scan')                           
                     else:
                             isError = 1        
                             status_info_widget.update(value='Nothing to scan') 
 
+
+                            
+
                 if not isError:
-                    scanned_files_text_widget.update(value=f'{SCANNED_FILES} {scan_result.scanned_files}') 
-                    infected_files_text_widget.update(value=f'{INFECTED_FILES} {scan_result.infected_files}') 
-                    dangerous_globals_text_widget.update(value=f'{DANGEROUS_GLOBALS} {scan_result.issues_count}') 
-                    scanned_files_len = len(scanned_files)
-                    infected_files_len = len(infected_files)
-
-                    print('')
-                    if scanned_files_len > 0:
-                        print('----------- SCANNED FILES -----------')
-                        for file in scanned_files:
-                            print(f"    {file}")
-
-                        if infected_files_len == 0:
-                            status_state_text_widget.update(text_color='white')
-                            status_state_text_widget.Widget.config(background=COLOR_DARK_GREEN)  
-                            status_state_text_widget.ParentRowFrame.config(background=COLOR_DARK_GREEN)
-                            status_state_text_widget.update(value='PASS') 
+                    if scan_result:
+                        scanned_files_text_widget.update(value=f'{SCANNED_FILES} {scan_result.scanned_files}') 
+                        infected_files_text_widget.update(value=f'{INFECTED_FILES} {scan_result.infected_files}') 
+                        dangerous_globals_text_widget.update(value=f'{DANGEROUS_GLOBALS} {scan_result.issues_count}') 
+                        scanned_files_len = len(scanned_files)
+                        infected_files_len = len(infected_files)
 
-                    if infected_files_len > 0:
                         print('')
-                        print('----------- INFECTED FILES -----------')
-                        for file in infected_files:
-                            print(f"    {file}")
+                        if scanned_files_len > 0:
+                            # print('----------- SCANNED FILES -----------')
+                            # for file in scanned_files:
+                            #     print(f"    {file}")
+
+                            if infected_files_len == 0:
+                                status_state_text_widget.update(text_color='white')
+                                status_state_text_widget.Widget.config(background=COLOR_DARK_GREEN)  
+                                status_state_text_widget.ParentRowFrame.config(background=COLOR_DARK_GREEN)
+                                status_state_text_widget.update(value='PASS') 
+
+                        if infected_files_len > 0:
+                            print('')
+                            print('----------- INFECTED FILES -----------')
+                            print('')
 
-                        status_state_text_widget.update(text_color='white')
-                        status_state_text_widget.Widget.config(background=COLOR_RED_ORANGE)  
-                        status_state_text_widget.ParentRowFrame.config(background=COLOR_RED_ORANGE)
-                        status_state_text_widget.update(value='INFECTED') 
+                            for file in infected_files:
+                                print(f"    {file}")
+
+                            status_state_text_widget.update(text_color='white')
+                            status_state_text_widget.Widget.config(background=COLOR_RED_ORANGE)  
+                            status_state_text_widget.ParentRowFrame.config(background=COLOR_RED_ORANGE)
+                            status_state_text_widget.update(value='FAIL') 
 
-                    status_info_widget.update(value='Done') 
+                        status_info_widget.update(value='Done') 
 
         if event == '-huggingface_radio-':
                 window['-frame_file-'].update(visible=False)  
diff --git a/picklescan/src/picklescan/scanner.py b/picklescan/src/picklescan/scanner.py
@@ -41,10 +41,13 @@ class ScanResult:
     infected_files: int = 0
 
     def merge(self, sr: "ScanResult"):
-        self.globals.extend(sr.globals)
-        self.scanned_files += sr.scanned_files
-        self.issues_count += sr.issues_count
-        self.infected_files += sr.infected_files
+            if sr:
+                self.globals.extend(sr.globals)
+                self.scanned_files += sr.scanned_files
+                self.issues_count += sr.issues_count
+                self.infected_files += sr.infected_files
+            else:
+                return None
 
 _log = logging.getLogger("picklescan")
 
@@ -136,13 +139,11 @@ def _http_get(url) -> bytes:
     _log.debug(f"Request: GET {url}")
 
     parsed_url = urllib.parse.urlparse(url)
-    # print('parsed_url',parsed_url.path)
     path_and_query = parsed_url.path + (
         "?" + parsed_url.query if len(parsed_url.query) > 0 else ""
     )
 
     conn = http.client.HTTPSConnection(parsed_url.netloc)
-    # print('path_and_query',path_and_query)
 
     try:
         conn.request("GET", path_and_query)
@@ -225,12 +226,10 @@ def scan_pickle_bytes(data: IO[bytes], file_id) -> ScanResult:
         if unsafe_filter is not None and (
             unsafe_filter == "*" or g.name in unsafe_filter
         ):
-            # print('----------- INFECTED FILE -----------')
             g.safety = SafetyLevel.Dangerous
             _log.warning(
                 "%s: %s import '%s %s' FOUND", file_id, g.safety.value, g.module, g.name
             )
-            # print(f"{file_id} {g.safety.value} import {g.module} {g.name} FOUND")
             infected_files.append(f"{file_id} {g.safety.value} import {g.module} {g.name} FOUND")
             issues_count += 1
         elif safe_filter is not None and (safe_filter == "*" or g.name in safe_filter):
@@ -244,21 +243,27 @@ def scan_pickle_bytes(data: IO[bytes], file_id) -> ScanResult:
 
 def scan_zip_bytes(data: IO[bytes], file_id) -> ScanResult:
     result = ScanResult([])
+    try:
+        with zipfile.ZipFile(data, "r") as zip:
+            file_names = zip.namelist()
+            _log.debug("Files in archive %s: %s", file_id, file_names)
+            for file_name in file_names:
+                if os.path.splitext(file_name)[1] in _pickle_file_extensions:
+                    _log.debug("Scanning file %s in zip archive %s", file_name, file_id)
+                    with zip.open(file_name, "r") as file:
+                        result.merge(scan_pickle_bytes(file, f"{file_id}:{file_name}"))
+        if result:
+            return result
+    except zipfile.BadZipFile as e:
+        print("  FAIL TO SCAN: Corrupted file",file_id)   
 
-    with zipfile.ZipFile(data, "r") as zip:
-        file_names = zip.namelist()
-        _log.debug("Files in archive %s: %s", file_id, file_names)
-        for file_name in file_names:
-            if os.path.splitext(file_name)[1] in _pickle_file_extensions:
-                _log.debug("Scanning file %s in zip archive %s", file_name, file_id)
-                with zip.open(file_name, "r") as file:
-                    result.merge(scan_pickle_bytes(file, f"{file_id}:{file_name}"))
 
-    return result
 
 
 def scan_pytorch(data: IO[bytes], file_id) -> ScanResult:
     # new pytorch format
+    print("  Scanning file:",file_id)
+
     if _is_zipfile(data):
         return scan_zip_bytes(data, file_id)
     # old pytorch format
@@ -273,7 +278,6 @@ def scan_pytorch(data: IO[bytes], file_id) -> ScanResult:
             except TarError:
                 # file does not contain a tar
                 data.seek(0)
-
         magic = get_magic_number(data)
         if magic != MAGIC_NUMBER:
             raise InvalidMagicError(magic, MAGIC_NUMBER)
@@ -286,14 +290,20 @@ def scan_pytorch(data: IO[bytes], file_id) -> ScanResult:
         for _ in range(5):
             scan_result.merge(scan_pickle_bytes(data, file_id))
         scan_result.scanned_files = 1
-        return scan_result
 
+        return scan_result
 
 def scan_bytes(data: IO[bytes], file_id, file_ext: Optional[str] = None) -> ScanResult:
     if file_ext is not None and file_ext in _pytorch_file_extensions:
-        return scan_pytorch(data, file_id)
+        try: 
+            return scan_pytorch(data, file_id)
+        except:
+            print("  FAIL TO SCAN: Corrupted file", file_id)
+            return None
     else:
         is_zip = zipfile.is_zipfile(data)
+        if isinstance(data,io.BufferedReader):
+            print("  Scanning file:",data.name)
         data.seek(0)
         return (
             scan_zip_bytes(data, file_id)
@@ -362,4 +372,5 @@ def scan_file_path(path) -> ScanResult:
 
 def scan_url(url) -> ScanResult:
     scanned_files.append(url)
+    print("  Scanning file: ",url)
     return scan_bytes(io.BytesIO(_http_get(url)), url)
