Implemented automatic scaling units for byte values

wordfence/cli/malwarescan/progress.py

@@ -12,6 +12,7 @@
 from ..banner.banner import get_welcome_banner
 from ...util import timing
 from ...util.unicode import filter_control_characters
+from ...util.units import scale_byte_unit
 
 
 class ProgressException(Exception):
@@ -507,12 +508,14 @@ def _get_metrics(
             ) -> List[Metric]:
         file_count = update.metrics.get_int_metric('counts', worker_index)
         byte_count = update.metrics.get_int_metric('bytes', worker_index)
+        byte_value = scale_byte_unit(byte_count)
         match_count = update.metrics.get_int_metric('matches', worker_index)
         file_rate = self._compute_rate(file_count, update.elapsed_time)
         byte_rate = self._compute_rate(byte_count, update.elapsed_time)
+        byte_rate = scale_byte_unit(byte_rate)
         metrics = [
                 Metric('Files Processed', file_count),
-                Metric('Bytes Processed', byte_count),
+                Metric('Bytes Processed', byte_value),
                 Metric('Matches Found', match_count),
                 Metric('Files / Second', file_rate),
                 Metric('Bytes / Second', byte_rate)

wordfence/scanning/scanner.py

@@ -15,6 +15,7 @@
 from ..util import timing
 from ..util.io import StreamReader
 from ..util.pcre import PcreOptions, PCRE_DEFAULT_OPTIONS, PcreJitStack
+from ..util.units import scale_byte_unit
 from ..intel.signatures import SignatureSet
 from ..logging import log, remove_initial_handler
 
@@ -501,15 +502,15 @@ def get_scan_finished_messages(
         ) -> ScanFinishedMessages:
     match_count = metrics.get_total_matches()
     total_count = metrics.get_total_count()
-    byte_count = metrics.get_total_bytes()
+    byte_value = scale_byte_unit(metrics.get_total_bytes())
     elapsed_time = round(timer.get_elapsed())
     timeout_count = metrics.get_total_timeouts()
     timeouts_message = None
     if timeout_count > 0:
         timeouts_message = f'{timeout_count} timeout(s) occurred during scan'
-    results_message = (f'Found {match_count} matching file(s) after '
+    results_message = (f'Found {match_count} suspicious file(s) after '
                        f'processing {total_count} file(s) containing '
-                       f'{byte_count} byte(s) over {elapsed_time} second(s)')
+                       f'{byte_value} over {elapsed_time} second(s)')
 
     if metrics.skipped_files > 0:
         skipped_message = (

wordfence/util/units.py

@@ -1,4 +1,6 @@
 import re
+from dataclasses import dataclass
+from enum import Enum
 
 KIBIBYTE = 1024
 MEBIBYTE = 1024 * 1024
@@ -23,3 +25,45 @@ def byte_length(conversion_value: str) -> int:
     if not sizings_map.get(suffix, False):
         raise ValueError("Unrecognized byte length suffix")
     return int(match.group(1)) * sizings_map.get(suffix)
+
+
+class ByteUnit(Enum):
+    BYTE = (1, 'B')
+    KIBIBYTE = (pow(2, 10), 'KiB')
+    MEBIBYTE = (pow(2, 20), 'MiB')
+    GIBIBYTE = (pow(2, 30), 'GiB')
+    TEBIBYTE = (pow(2, 40), 'TiB')
+
+    def __init__(
+                self,
+                size: int,
+                abbreviation: str
+            ):
+        self.size = size
+        self.abbreviation = abbreviation
+
+
+@dataclass
+class ByteUnitValue:
+    value: float
+    unit: ByteUnit
+
+    def __str__(self) -> str:
+        if self.unit.size == 1:
+            value = int(self.value)
+        else:
+            value = round(self.value, 1)
+        return f'{value} {self.unit.abbreviation}'
+
+
+def scale_byte_unit(byte_count: int) -> ByteUnitValue:
+    scaled_unit = ByteUnit.BYTE
+    for unit in ByteUnit:
+        if byte_count >= unit.size \
+                and (
+                    scaled_unit is None
+                    or unit.size > scaled_unit.size
+                ):
+            scaled_unit = unit
+    scaled_value = byte_count / scaled_unit.size
+    return ByteUnitValue(scaled_value, scaled_unit)