utils/color_functions: color only inside atty

Before this patch, piping rainbow output to a file would keep
ANSI escape codes, making it harder to read and parse the output.

This patch moves pygments highlight function to utils.color_functions
and then configures it to only output ANSI colors if isatty is true.

Author: erdnaxe

rainbow/rainbow.py

@@ -24,15 +24,11 @@
 from typing import Callable, Tuple, Optional, List, Dict, Set, Any
 import capstone as cs
 import unicorn as uc
-from pygments import highlight
-from pygments.formatters.terminal import TerminalFormatter
-from pygments.lexers.asm import NasmLexer
-# TODO: Add note about colorama use. Call init in example code.
 from unicorn import UcError
 
 from .leakage_models import LeakageModel
 from .utils import region_intersects, HookWeakMethod
-from .utils.color_functions import color
+from .utils.color_functions import color, highlight_asmline
 from .loaders import load_selector
 
 
@@ -117,10 +113,6 @@ def __init__(self, print_config: Print = Print(0), trace_config: TraceConfig = T
         self.last_address = 0
         self.trace = []
 
-        # Prepare the formatters
-        self.asm_hl = NasmLexer()
-        self.asm_fmt = TerminalFormatter(outencoding="utf-8")
-
         # Prepare the emulator and disassembler
         self.emu = uc.Uc(self.UC_ARCH, self.UC_MODE)
         self.disasm = cs.Cs(self.CS_ARCH, self.CS_MODE)
@@ -402,15 +394,6 @@ def disassemble_single_detailed(self, addr: int, size: int) -> cs.CsInsn:
         insn = self.emu.mem_read(addr, 2 * size)
         return self._disassemble_cache(self.disasm.disasm, bytes(insn), addr)
 
-    def print_asmline(self, adr, ins, op_str):
-        """ Pretty-print assembly using pygments syntax highlighting """
-        line = (
-            highlight(f"{ins:<6}  {op_str:<20}", self.asm_hl, self.asm_fmt)
-                .decode()
-                .strip("\n")
-        )
-        print("\n" + color("YELLOW", f"{adr:8X}  ") + line, end=";")
-
     def hook_prolog(self, name, fn):
         """
         Add a call to function 'fn' when 'name' is called during execution.
@@ -569,9 +552,9 @@ def _code_hook(self, uci, address, size, _):
         if self.print_config & Print.Code:
             if ins is None:
                 adr, size, _ins, op_str = self.disassemble_single(address, size)
-                self.print_asmline(adr, _ins, op_str)
+                highlight_asmline(adr, _ins, op_str)
             else:
-                self.print_asmline(address, ins.mnemonic, ins.op_str)
+                highlight_asmline(address, ins.mnemonic, ins.op_str)
 
         # Handle the register tracing
         event = None

rainbow/utils/color_functions.py

@@ -15,21 +15,39 @@
 #
 #
 # Copyright 2019 Victor Servant, Ledger SAS
+import sys
 
-_RST = "\u001b[0m"
-FOREGROUND_COLORS = {
-    "BLACK": "\u001b[30m",
-    "RED": "\u001b[31m",
-    "GREEN": "\u001b[32m",
-    "YELLOW": "\u001b[33m",
-    "BLUE": "\u001b[34m",
-    "MAGENTA": "\u001b[35m",
-    "CYAN": "\u001b[36m",
-    "WHITE": "\u001b[37m",
-}
+from pygments import highlight
+from pygments.formatters import NullFormatter, TerminalFormatter
+from pygments.lexers.asm import NasmLexer
+
+# Use colors only if tty support ANSI colors
+ASM_HL = NasmLexer()
+ASM_FMT = NullFormatter(outencoding="utf-8")
+FOREGROUND_COLORS = {}
+if sys.stdout.isatty():
+    ASM_FMT = TerminalFormatter(outencoding="utf-8")
+    FOREGROUND_COLORS = {
+        "BLACK": "\u001b[30m",
+        "RED": "\u001b[31m",
+        "GREEN": "\u001b[32m",
+        "YELLOW": "\u001b[33m",
+        "BLUE": "\u001b[34m",
+        "MAGENTA": "\u001b[35m",
+        "CYAN": "\u001b[36m",
+        "WHITE": "\u001b[37m",
+        "_RST": "\u001b[0m",
+    }
 
 
 def color(color_name: str, x: str) -> str:
     """Color string `x` with color `color_name`"""
     color_code = FOREGROUND_COLORS.get(color_name, "")
-    return f"{color_code}{x}{_RST}"
+    rst_code = FOREGROUND_COLORS.get("_RST", "")
+    return f"{color_code}{x}{rst_code}"
+
+
+def highlight_asmline(addr: int, ins: str, op_str: str):
+    """Pretty-print assembly using pygments syntax highlighting"""
+    line = highlight(f"{ins:<6}  {op_str:<20}", ASM_HL, ASM_FMT).decode().strip("\n")
+    print("\n" + color("YELLOW", f"{addr:8X}  ") + line, end=";")