From 491e2b5462dead414839fc32b5a540bc921d4d34 Mon Sep 17 00:00:00 2001
From: Leo Feradero Nugraha <leoferaderonugraha@gmail.com>
Date: Mon, 6 Feb 2023 05:04:46 +0700
Subject: [PATCH] Fix improper appended byte, add support for python3

---
 tools/pump.py | 48 ++++++++++++++++++++++++++++++++----------------
 1 file changed, 32 insertions(+), 16 deletions(-)

diff --git a/tools/pump.py b/tools/pump.py
index 61030cfa..2fdf3be6 100755
--- a/tools/pump.py
+++ b/tools/pump.py
@@ -1,25 +1,41 @@
+"""
+File pumper. Increase file size with null byte(s) at the end of the file.
+"""
+
 import sys
-#python fpump.py [file] [size] [-mb/-kb]
+import os
+
+# python pump.py [file] [size] [-mb/-kb]
+
+# Also refer to KiB/MiB
+KB = 1024
+MB = KB * 1024
 
 if len(sys.argv) < 4:
- sys.exit('[-] Missing argument!\n[+] Usage: python pumper.py [file] [size] [-mb/-kb]')
+    sys.stderr.write('[-] Missing argument!\n')
+    sys.stderr.write('[+] Usage: python pumper.py [file] [size] [-mb/-kb]\n')
+    exit(1)
 
-fp = sys.argv[1]
+fileName = sys.argv[1]
 size = int(sys.argv[2])
-tp = sys.argv[3]
+unit = sys.argv[3]
+
+if not os.path.exists(fileName):
+    sys.stderr.write('[-] File {!r} is not exists!\n'.format(fileName))
+    exit(1)
 
-f = open(fp, 'ab')
-if tp == '-kb':
-    b_size = size * 1024
-elif tp == '-mb':
-    b_size = size * 1048576
-else:
-    sys.exit('[-] Use -mb or -kb!')
+if unit != '-mb' and unit != '-kb':
+    sys.stderr.write('[-] Use -mb or -kb!\n')
+    exit(1)
 
-bufferSize = 256
-for i in range(b_size/bufferSize):
-    f.write(str('0' * bufferSize))
+with open(fileName, 'ab') as fp:
+    if unit == '-kb':
+        blockSize = size * KB
+    elif unit == '-mb':
+        blockSize = size * MB
 
-f.close()
+    bufferSize = 256
+    for i in range(blockSize // bufferSize):
+        fp.write(('\0' * bufferSize).encode('utf-8'))
 
-print '[+] Finished pumping', fp, 'with', size, tp
+print('[+] Finished pumping {!r} with {}{}'.format(fileName, size, unit[1:]))