Add redirect print to capa main

CHANGELOG.md

@@ -22,6 +22,7 @@
 - improve ELF strtab and needed parsing @mr-tz
 - better handle exceptional cases when parsing ELF files [#1458](https://github.com/mandiant/capa/issues/1458) [@Aayush-Goel-04](https://github.com/aayush-goel-04)
 - Improved testing coverage for Binary Ninja Backend [#1446](https://github.com/mandiant/capa/issues/1446) [@Aayush-Goel-04](https://github.com/aayush-goel-04)
+- Add redirect print to tqdm for capa main [#749](https://github.com/mandiant/capa/issues/749) [@Aayush-Goel-04](https://github.com/aayush-goel-04)
 - extractor: fix binja installation path detection does not work with Python 3.11
 
 ### capa explorer IDA Pro plugin

capa/helpers.py

@@ -6,9 +6,13 @@
 #  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and limitations under the License.
 import os
+import inspect
 import logging
+import contextlib
 from typing import NoReturn
 
+import tqdm
+
 from capa.exceptions import UnsupportedFormatError
 from capa.features.common import FORMAT_PE, FORMAT_SC32, FORMAT_SC64, FORMAT_DOTNET, FORMAT_UNKNOWN, Format
 
@@ -85,6 +89,39 @@ def get_format(sample: str) -> str:
     return FORMAT_UNKNOWN
 
 
+@contextlib.contextmanager
+def redirecting_print_to_tqdm(disable_progress):
+    """
+    tqdm (progress bar) expects to have fairly tight control over console output.
+    so calls to `print()` will break the progress bar and make things look bad.
+    so, this context manager temporarily replaces the `print` implementation
+    with one that is compatible with tqdm.
+    via: https://stackoverflow.com/a/42424890/87207
+    """
+    old_print = print
+
+    def new_print(*args, **kwargs):
+        # If tqdm.tqdm.write raises error, use builtin print
+        if disable_progress:
+            old_print(*args, **kwargs)
+        else:
+            try:
+                tqdm.tqdm.write(*args, **kwargs)
+            except:
+                old_print(*args, **kwargs)
+
+    try:
+        # Globally replace print with new_print.
+        # Verified this works manually on Python 3.11:
+        #     >>> import inspect
+        #     >>> inspect.builtins
+        #     <module 'builtins' (built-in)>
+        inspect.builtins.print = new_print  # type: ignore
+        yield
+    finally:
+        inspect.builtins.print = old_print  # type: ignore
+
+
 def log_unsupported_format_error():
     logger.error("-" * 80)
     logger.error(" Input file does not appear to be a PE or ELF file.")

capa/main.py

@@ -53,6 +53,7 @@ from capa.helpers import (
     get_file_taste,
     get_auto_format,
     log_unsupported_os_error,
+    redirecting_print_to_tqdm,
     log_unsupported_arch_error,
     log_unsupported_format_error,
 )
@@ -251,6 +252,7 @@ def find_capabilities(ruleset: RuleSet, extractor: FeatureExtractor, disable_pro
         "library_functions": {},
     }  # type: Dict[str, Any]
 
+    with redirecting_print_to_tqdm(disable_progress):
         pbar = tqdm.tqdm
         if disable_progress:
             # do not use tqdm to avoid unnecessary side effects when caller intends

scripts/lint.py

@@ -22,13 +22,11 @@ import time
 import string
 import difflib
 import hashlib
-import inspect
 import logging
 import pathlib
 import argparse
 import itertools
 import posixpath
-import contextlib
 from typing import Set, Dict, List
 from pathlib import Path
 from dataclasses import field, dataclass
@@ -866,37 +864,6 @@ def width(s, count):
         return s.ljust(count)
 
 
-@contextlib.contextmanager
-def redirecting_print_to_tqdm():
-    """
-    tqdm (progress bar) expects to have fairly tight control over console output.
-    so calls to `print()` will break the progress bar and make things look bad.
-    so, this context manager temporarily replaces the `print` implementation
-    with one that is compatible with tqdm.
-
-    via: https://stackoverflow.com/a/42424890/87207
-    """
-    old_print = print
-
-    def new_print(*args, **kwargs):
-        # If tqdm.tqdm.write raises error, use builtin print
-        try:
-            tqdm.tqdm.write(*args, **kwargs)
-        except:
-            old_print(*args, **kwargs)
-
-    try:
-        # Globally replace print with new_print.
-        # Verified this works manually on Python 3.11:
-        #     >>> import inspect
-        #     >>> inspect.builtins
-        #     <module 'builtins' (built-in)>
-        inspect.builtins.print = new_print  # type: ignore
-        yield
-    finally:
-        inspect.builtins.print = old_print  # type: ignore
-
-
 def lint(ctx: Context):
     """
     Returns: Dict[string, Tuple(int, int)]
@@ -907,7 +874,7 @@ def lint(ctx: Context):
 
     source_rules = [rule for rule in ctx.rules.rules.values() if not rule.is_subscope_rule()]
     with tqdm.contrib.logging.tqdm_logging_redirect(source_rules, unit="rule") as pbar:
-        with redirecting_print_to_tqdm():
+        with capa.helpers.redirecting_print_to_tqdm(False):
             for rule in pbar:
                 name = rule.name
                 pbar.set_description(width(f"linting rule: {name}", 48))