From 28ad261c83b6fb426caf1b8b4a861fcb5025acf2 Mon Sep 17 00:00:00 2001
From: Cyberes <cyberes@evulid.cc>
Date: Sat, 11 Nov 2023 12:06:32 -0700
Subject: [PATCH] adjust downloader and conversion

---
 exfiltrate.py |  4 ----
 pkg/image.py  | 17 +++++++++++------
 pkg/thread.py | 24 ++++++++++++++----------
 3 files changed, 25 insertions(+), 20 deletions(-)

diff --git a/exfiltrate.py b/exfiltrate.py
index 0bedaea..912efb8 100644
--- a/exfiltrate.py
+++ b/exfiltrate.py
@@ -107,16 +107,12 @@ if __name__ == '__main__':
                                 tiles.add((result_row, result_col))
                             elif new_image == 'failure':
                                 retry_files.add((result_row, result_col))
-                                update_bar_postfix()
                             elif new_image == 'fixed':
                                 tiles.add((result_row, result_col))
                                 fixed_files += 1
-                                update_bar_postfix()
                             elif new_image == 'converted':
                                 tiles.add((result_row, result_col))
                                 converted_files += 1
-                                # Don't update bar on every loop since we could be converting many files.
-                                # update_bar_postfix()
                             col_bar.update()
                             row_bar.refresh()
                 col_bar.close()
diff --git a/pkg/image.py b/pkg/image.py
index 3055f79..9a1c2d4 100644
--- a/pkg/image.py
+++ b/pkg/image.py
@@ -29,9 +29,14 @@ def is_png(file_path):
 
 
 def convert_to_png(file_path):
-    img = Image.open(file_path)
-    if img.format != 'PNG':
-        img.save(file_path, format='PNG')
-        # tqdm.write(f'Converted {file_path} from {img.format} to PNG')
-        return True
-    return False
+    try:
+        img = Image.open(file_path)
+        if img.format != 'PNG':
+            img.save(file_path, format='PNG')
+            return True
+        else:
+            return False
+    except Exception:
+        # import traceback
+        # traceback.print_exc()
+        return False
diff --git a/pkg/thread.py b/pkg/thread.py
index ce1c159..f12e6ba 100644
--- a/pkg/thread.py
+++ b/pkg/thread.py
@@ -40,22 +40,26 @@ def download_tile(task):
         tile_url = f"{base_url}/{row}/{col}".replace('//', '/').replace(':/', '://')
         response = requests.get(tile_url, headers=r_headers, proxies=PROXIES if use_proxy else None, timeout=60)
         if response.status_code == 200:
-            if not do_convert_to_png and not response.headers.get('Content-Type') == 'image/png':
-                # If we will convert the image to a PNG, ignore this header.
-                raise Exception(f'Response gave Content-Type: {response.headers.get("Content-Type")}')
+            # if not do_convert_to_png and not response.headers.get('Content-Type') == 'image/png':
+            #     # If we will convert the image to a PNG, ignore this header.
+            #     raise Exception(f'Response gave Content-Type: {response.headers.get("Content-Type")}')
             with open(output_path, "wb") as f:
                 f.write(response.content)
 
             if do_convert_to_png:
                 convert_to_png(output_path)
-                corrupted_image = True  # force a re-check
-
-            # Recheck the PNG if it was corrupted.
-            if corrupted_image and not is_png(output_path):
-                print(f"Retry for {row}_{col} failed a second time: cannot identify image file")
-                return row, col, 'failure'
+                if not is_png(output_path)[0]:
+                    tqdm.write(f'PNG conversion for {output_path} failed')
+                else:
+                    return row, col, 'success' if not corrupted_image else 'fixed'
             else:
-                return row, col, 'success' if not corrupted_image else 'fixed'
+                # Recheck the PNG if it was corrupted.
+                valid_png_file, image_type = is_png(output_path)
+                if not valid_png_file:
+                    tqdm.write(f'Bad image file: "{output_path}" (is format: {image_type}).')
+                    return row, col, 'failure'
+                else:
+                    return row, col, 'success' if not corrupted_image else 'fixed'
         else:
             print(f"Failed to download tile {row}_{col}")
             return row, col, 'failure'