From 3fb31cc42cf99c0a9e44ed55929edf0489004cd3 Mon Sep 17 00:00:00 2001
From: Yun Zheng Hu <hu@fox-it.com>
Date: Wed, 15 Dec 2021 10:49:23 +0100
Subject: [PATCH] Output scanning stats and version information

Version information can be suppressed using -b or --no-banner
---
 log4j-finder.py | 61 ++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 48 insertions(+), 13 deletions(-)

diff --git a/log4j-finder.py b/log4j-finder.py
index 54263c5..07f1eaf 100755
--- a/log4j-finder.py
+++ b/log4j-finder.py
@@ -18,6 +18,7 @@
 import os
 import io
 import sys
+import time
 import zipfile
 import logging
 import argparse
@@ -29,13 +30,23 @@
 
 from pathlib import Path
 
+__version__ = "1.0.1"
+FIGLET = f"""\
+ __               _____  __         ___ __           __
+|  |.-----.-----.|  |  ||__|______.'  _|__|.-----.--|  |.-----.----.
+|  ||  _  |  _  ||__    |  |______|   _|  ||     |  _  ||  -__|   _|
+|__||_____|___  |   |__||  |      |__| |__||__|__|_____||_____|__|
+          |_____|      |___| v{__version__} https://github.com/fox-it/log4j-finder
+"""
+
 # Optionally import colorama to enable colored output for Windows
 try:
     import colorama
+
     colorama.init()
     NO_COLOR = False
 except ImportError:
-    NO_COLOR = True if sys.platform == 'win32' else False
+    NO_COLOR = True if sys.platform == "win32" else False
 
 log = logging.getLogger(__name__)
 
@@ -82,15 +93,17 @@ def md5_digest(fobj):
     return d.hexdigest()
 
 
-def iter_scandir(path):
+def iter_scandir(path, stats=None):
     """
     Yields all files matcthing JAR_EXTENSIONS or FILENAMES recursively in path
     """
     p = Path(path)
     if p.is_file():
+        if stats:
+            stats["files"] += 1
         yield p
     try:
-        for entry in scantree(path):
+        for entry in scantree(path, stats=stats):
             if entry.is_symlink():
                 continue
             elif entry.is_file():
@@ -103,20 +116,24 @@ def iter_scandir(path):
         log.debug(e)
 
 
-def scantree(path):
+def scantree(path, stats=None):
     """Recursively yield DirEntry objects for given directory."""
     try:
         with os.scandir(path) as it:
             for entry in it:
                 if entry.is_dir(follow_symlinks=False):
-                    yield from scantree(entry.path)
+                    if stats:
+                        stats["directories"] += 1
+                    yield from scantree(entry.path, stats=stats)
                 else:
+                    if stats:
+                        stats["files"] += 1
                     yield entry
     except IOError as e:
         log.debug(e)
 
 
-def iter_jarfile(fobj, parents=None):
+def iter_jarfile(fobj, parents=None, stats=None):
     """
     Yields (zfile, zinfo, zpath, parents) for each file in zipfile that matches `FILENAMES` or `JAR_EXTENSIONS` (recursively)
     """
@@ -133,9 +150,9 @@ def iter_jarfile(fobj, parents=None):
                         zfile.open(zinfo.filename), parents=parents + [zpath]
                     )
     except IOError as e:
-        log.debug(f"{fobj}: {e}", e)
+        log.debug(f"{fobj}: {e}")
     except zipfile.BadZipFile as e:
-        log.debug(f"{fobj}: {e}", e)
+        log.debug(f"{fobj}: {e}")
 
 
 def red(s):
@@ -210,6 +227,7 @@ def main():
     parser.add_argument(
         "-n", "--no-color", action="store_true", help="disable color output"
     )
+    parser.add_argument("-b", "--no-banner", action="store_true", help="disable banner")
     args = parser.parse_args()
     logging.basicConfig(
         format="%(asctime)s %(levelname)s %(message)s",
@@ -226,20 +244,29 @@ def main():
         NO_COLOR = True
 
     stats = {
+        "scanned": 0,
+        "files": 0,
+        "directories": 0,
         "vulnerable": 0,
         "good": 0,
         "unknown": 0,
     }
+    start_time = time.monotonic()
+
+    if not args.no_banner:
+        print(FIGLET)
     for directory in args.path:
-        print("Scanning:", directory)
-        for p in iter_scandir(directory):
+        print(f"[{datetime.datetime.utcnow()}] Scanning: {directory}")
+        for p in iter_scandir(directory, stats=stats):
             if p.name.lower() in FILENAMES:
+                stats["scanned"] += 1
                 log.info(f"Found file: {p}")
                 with p.open("rb") as fobj:
                     check_vulnerable(fobj, [p], stats)
             if p.suffix.lower() in JAR_EXTENSIONS:
                 try:
                     log.info(f"Found jar file: {p}")
+                    stats["scanned"] += 1
                     for (zinfo, zfile, zpath, parents) in iter_jarfile(
                         p.resolve().open("rb"), parents=[p.resolve()]
                     ):
@@ -249,13 +276,21 @@ def main():
                 except IOError as e:
                     log.debug(f"{p}: {e}", e)
 
+    elapsed_time = time.monotonic() - start_time
+    print(
+        f"[{datetime.datetime.utcnow()}] Finished scan, elapsed time: {elapsed_time:.2f} seconds"
+    )
+
     print("\nSummary:")
+    print(f" Processed {stats['files']} files and {stats['directories']} directories")
+    print(f" Scanned {stats['scanned']} files")
     if stats["vulnerable"]:
-        print(" Found {} vulnerable files".format(stats["vulnerable"]))
+        print("  Found {} vulnerable files".format(stats["vulnerable"]))
     if stats["good"]:
-        print(" Found {} good files".format(stats["good"]))
+        print("  Found {} good files".format(stats["good"]))
     if stats["unknown"]:
-        print(" Found {} unknown files".format(stats["unknown"]))
+        print("  Found {} unknown files".format(stats["unknown"]))
+    print(f"\nElapsed time: {elapsed_time:.2f} seconds ")
 
 
 if __name__ == "__main__":