improve mass_whitelist.py output

1 files changed, 14 insertions, 4 deletions

diff --git a/php-malware-finder/utils/mass_whitelist.py b/php-malware-finder/utils/mass_whitelist.py
index 82733fa..ef7e39d 100755
--- a/php-malware-finder/utils/mass_whitelist.py
+++ b/php-malware-finder/utils/mass_whitelist.py
@@ -1,6 +1,8 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
+from __future__ import print_function
+
 import sys
 import tarfile
 from copy import copy
@@ -65,7 +67,8 @@ def generate_whitelist(version):
     dl_failed = False
     download_url = Opts.URL_PATTERN.replace('__version__', version)
     download_url_str = Opts.URL_PATTERN.replace('__version__', '\x1b[1;33m%s\x1b[0m' % version)
-    print "[+] Downloading %s... " % download_url_str,
+    print("[+] Downloading %s... " % download_url_str, end='')
+    sys.stdout.flush()
     try:
         resp = urlopen(download_url)
         resp_code = resp.code
@@ -73,13 +76,15 @@ def generate_whitelist(version):
         dl_failed = True
         resp_code = err.code
     if dl_failed or (resp_code != 200):
-        print "\x1b[1;31mFAILED (%d)\x1b[0m" % resp_code
+        print("\x1b[1;31mFAILED (%d)\x1b[0m" % resp_code)
         return None
     data = StringIO(resp.read())
     data.seek(0)
-    print "\x1b[1;32mOK\x1b[0m"
+    print("\x1b[1;32mOK\x1b[0m")
 
     # extract archive and check against YARA signatures (in-memory)
+    print("[-] Generating whitelist... ", end='')
+    sys.stdout.flush()
     tar = tarfile.open(mode='r:gz', fileobj=data)
     for entry in tar.getnames():
         entry_fd = tar.extractfile(entry)
@@ -89,6 +94,7 @@ def generate_whitelist(version):
         matches = Opts.YARA_RULES.match(data=entry_data, fast=True)
         if matches:
             rules['/'.join(entry.split('/')[1:])] = sha1(entry_data).hexdigest()
+    print("\x1b[1;32mDONE\x1b[0m")
 
     return rules
 
@@ -148,6 +154,7 @@ for vmajor in range(Opts.MIN_MAJOR, Opts.MAX_MAJOR + 1):
             first_ploop = False
 
 # remove duplicate entries:
+print("[+] Deduplicating detections... ", end='')
 known_files = []
 for version, rules in copy(whitelists.items()):
     used_rules = 0
@@ -160,7 +167,9 @@ for version, rules in copy(whitelists.items()):
             used_rules += 1
     if used_rules == 0:
         del whitelists[version]
+print("\x1b[1;32mDONE\x1b[0m")
 
+print("[+] Generating final whitelist... ", end='')
 # build final rule
 prefix = 8 * ' '
 conditions = []
@@ -174,6 +183,7 @@ for index, (version, rules) in enumerate(whitelists.items()):
         else:
             cond_str += '%shash.sha1(0, filesize) == "%s" or // %s\n' % (prefix, digest, filename)
     conditions.append(cond_str)
+print("\x1b[1;32mDONE\x1b[0m")
 
 final_rule = """
 import "hash"
@@ -191,4 +201,4 @@ private rule %(name)s
     'gendate': datetime.now().isoformat(),
     'conditions': '\n'.join(conditions)
 }
-print final_rule
+print(final_rule)