diff --git a/exfiltrate.py b/exfiltrate.py index 4d4f0f7..e3d2cd1 100644 --- a/exfiltrate.py +++ b/exfiltrate.py @@ -67,13 +67,21 @@ if __name__ == '__main__': row_bar = tqdm(total=0, desc='Row 000 | Loop 0/0', postfix={'new_files': 0, 'failures': 0, 'fixed': 0}) for i in range(1, args.download_loops + 1): row_bar.reset() + converted_files = 0 fixed_files = 0 total_downloaded = 0 row_i = min_row row_iter = range(min_row, max_row + 1) row_bar.total = len(row_iter) row_bar.desc = f'Row {row_i} | Loop {i}/{args.download_loops}' - row_bar.set_postfix({'new_files': total_downloaded, 'failures': len(retry_files), 'fixed': fixed_files}) + + + def update_bar_postfix(): + row_bar.set_postfix({'new': total_downloaded, 'failures': len(retry_files), 'fixed': fixed_files, 'converted': converted_files}) + + + update_bar_postfix() + for row in row_iter: row_i = row col_iter = range(min_col, max_col + 1) @@ -99,15 +107,19 @@ if __name__ == '__main__': tiles.add((result_row, result_col)) elif new_image == 'failure': retry_files.add((result_row, result_col)) - row_bar.set_postfix({'new_files': total_downloaded, 'failures': len(retry_files), 'fixed': fixed_files}) + update_bar_postfix() elif new_image == 'fixed': tiles.add((result_row, result_col)) fixed_files += 1 - row_bar.set_postfix({'new_files': total_downloaded, 'failures': len(retry_files), 'fixed': fixed_files}) + update_bar_postfix() + elif new_image == 'converted': + tiles.add((result_row, result_col)) + converted_files += 1 + update_bar_postfix() col_bar.update() row_bar.refresh() col_bar.close() - row_bar.set_postfix({'new_files': total_downloaded, 'failures': len(retry_files), 'fixed': fixed_files}) + update_bar_postfix() row_bar.update() row_bar.close() diff --git a/pkg/image.py b/pkg/image.py index 4538502..3055f79 100644 --- a/pkg/image.py +++ b/pkg/image.py @@ -28,8 +28,10 @@ def is_png(file_path): return False, None -def convert(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') + # tqdm.write(f'Converted {file_path} from {img.format} to PNG') + return True + return False diff --git a/pkg/thread.py b/pkg/thread.py index 5e041b5..ce1c159 100644 --- a/pkg/thread.py +++ b/pkg/thread.py @@ -5,7 +5,7 @@ import requests from tqdm import tqdm from pkg.proxies import PROXIES -from .image import convert, is_png +from .image import convert_to_png, is_png def del_path(p: Path): @@ -16,20 +16,20 @@ def del_path(p: Path): def download_tile(task): - row, col, base_url, r_headers, output, use_proxy, convert_to_png = task + row, col, base_url, r_headers, output, use_proxy, do_convert_to_png = task corrupted_image = False try: output_path: Path = output / f"{row}_{col}.png" if output_path.exists(): valid_png_file, image_type = is_png(output_path) - if convert_to_png and image_type and image_type != 'PNG': + if do_convert_to_png and image_type and image_type != 'PNG': # The image was read sucessfully by PIL but it's in the wrong format. - convert(output_path) + coverted = convert_to_png(output_path) if not is_png(output_path): tqdm.write(f'PNG conversion for {output_path} failed. Deleting and retrying...') corrupted_image = True else: - return row, col, 'exist' + return row, col, 'converted' if coverted else 'exist' elif not valid_png_file: # We will re-download the image. Don't need to delete it, just overwrite it. # del_path(output_path) @@ -40,14 +40,14 @@ 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 convert_to_png and not response.headers.get('Content-Type') == 'image/png': + 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 convert_to_png: - convert(output_path) + if do_convert_to_png: + convert_to_png(output_path) corrupted_image = True # force a re-check # Recheck the PNG if it was corrupted.