Move removing control characters to output.sanitize_text()

osc/core.py

@@ -53,6 +53,7 @@ from . import oscerr
 from . import output
 from . import store as osc_store
 from .connection import http_request, http_GET, http_POST, http_PUT, http_DELETE
+from .output import sanitize_text
 from .store import Store
 from .util import xdg
 from .util.helper import decode_list, decode_it, raw_input, _html_escape
@@ -6998,11 +6999,9 @@ def print_buildlog(
     def print_data(data, strip_time=False):
         if strip_time:
             data = buildlog_strip_time(data)
-        output_buffer.write(data.translate(all_bytes, remove_bytes))
+        # to protect us against control characters (CVE-2012-1095)
+        output_buffer.write(sanitize_text(data))
 
-    # to protect us against control characters (CVE-2012-1095)
-    all_bytes = bytes.maketrans(b'', b'')
-    remove_bytes = all_bytes[:8] + all_bytes[14:32]  # accept tabs and newlines
     query = {'nostream': '1', 'start': f'{offset}'}
     if last:
         query['last'] = 1

osc/output/__init__.py

@@ -1,6 +1,9 @@
 from .key_value_table import KeyValueTable
 from .input import get_user_input
 from .output import print_msg
+from .output import sanitize_text
+from .output import safe_print
+from .output import safe_write
 from .tty import colorize
 from .widechar import wc_ljust
 from .widechar import wc_width

osc/output/output.py

@@ -1,5 +1,8 @@
+import os
 import sys
 from typing import Optional
+from typing import TextIO
+from typing import Union
 
 from . import tty
 
@@ -39,3 +42,48 @@ def print_msg(*args, print_to: Optional[str] = "debug"):
         print(*args, file=sys.stderr)
     else:
         raise ValueError(f"Invalid value of the 'print_to' option: {print_to}")
+
+
+# Forbidden characters are nearly all control characters 0-31 with the exception of:
+#   0x09 - horizontal tab (\t)
+#   0x0A - line feed (\n)
+#   0x0D - carriage return (\r)
+# (related to CVE-2012-1095)
+#
+# It would be good to selectively allow 0x1B with safe & trusted escape sequences.
+FORBIDDEN_BYTES = b"\x00\x01\x02\x03\x04\x05\x06\x07\x08\x0b\x0c\x0e\x0f\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f"
+FORBIDDEN_CHARS = dict.fromkeys(FORBIDDEN_BYTES)
+
+
+def sanitize_text(text: Union[bytes, str]) -> Union[bytes, str]:
+    """
+    Remove forbidden characters from ``text``.
+    """
+    if isinstance(text, bytes):
+        return text.translate(None, FORBIDDEN_BYTES)
+    return text.translate(FORBIDDEN_CHARS)
+
+
+def safe_print(*args, **kwargs):
+    """
+    A wrapper to print() that runs sanitize_text() on all arguments.
+    """
+    args = [sanitize_text(i) for i in args]
+    print(*args, **kwargs)
+
+
+def safe_write(file: TextIO, text: Union[str, bytes], *, add_newline: bool = False):
+    """
+    Run sanitize_text() on ``text`` and write it to ``file``.
+
+    :param add_newline: Write a newline after writing the ``text``.
+    """
+    text = sanitize_text(text)
+    if isinstance(text, bytes):
+        file.buffer.write(text)
+        if add_newline:
+            file.buffer.write(os.linesep.encode("utf-8"))
+    else:
+        file.write(text)
+        if add_newline:
+            file.write(os.linesep)