retronews

commit df64ed1752e69b8a46c1fad85b898732b06d2138
parent 0a57a6665a732eaa6f82429acb453a5c48478668
Author: luke8086 <55237178+luke8086@users.noreply.github.com>
Date:   Wed, 14 Jun 2023 15:23:44 +0000

Improve code for sanitizing message text

Diffstat:
M
retronews.py
 | 
29
++++++++++++++++-------------
M
tests.py
 | 
2
+-

2 files changed, 17 insertions(+), 14 deletions(-)
diff --git a/retronews.py b/retronews.py
@@ -23,6 +23,7 @@ import logging
 import os
 import re
 import sqlite3
+import unicodedata
 import urllib.request
 from datetime import datetime
 from functools import partial, reduce
@@ -352,6 +353,18 @@ def wrap_paragraph(text: str) -> list[str]:
     return wrap(text, subsequent_indent=indent, break_on_hyphens=False, break_long_words=False)
 
 
+def sanitize_text(text: Optional[str]) -> str:
+    # For safety, remove any control characters except for \n and \t
+    # At least on HN some messages contain \x00 characters
+
+    text = text or ""
+    allowed_cc = set(("\n", "\t"))
+    chars = (c for c in text if c in allowed_cc or unicodedata.category(c) != "Cc")
+    text = "".join(chars)
+
+    return text
+
+
 def parse_html(html: str) -> list[str]:
     # This parser works well for HN messages because their markup is simple, and it can do
     # some custom optimizations, like expanding ellipsis-shortened links, preserving quote
@@ -655,13 +668,8 @@ def msg_flatten_thread(msg: Message, prefix: str = "", is_last_child: bool = Fal
             yield child
 
 
-def msg_sanitize_lines(lines: list[str]) -> list[str]:
-    # Some HN messages include null characters, which crash ncurses
-    return [line.replace("\u0000", "") for line in lines]
-
-
 def msg_build_raw_lines(msg: Message) -> list[str]:
-    text = msg.body or ""
+    text = sanitize_text(msg.body)
 
     # Unescape selected entities for better readability
     repl = {"/": "/", "'": "'", """: '"'}
@@ -680,16 +688,11 @@ def msg_build_lines(msg: Message) -> list[str]:
         "",
     ]
 
-    lines += parse_html(msg.body or "") if not msg.is_deleted else ["<deleted>"]
+    lines += parse_html(sanitize_text(msg.body)) if not msg.is_deleted else ["<deleted>"]
 
     return lines
 
 
-def msg_update_lines(msg: Message, raw_mode: bool = False) -> None:
-    msg.lines = msg_build_raw_lines(msg) if raw_mode else msg_build_lines(msg)
-    msg.lines = msg_sanitize_lines(msg.lines)
-
-
 def msg_unload(msg: Message) -> Message:
     msg.children = []
     msg.body = None
@@ -817,7 +820,7 @@ def app_refresh_message(app: AppState) -> None:
 
     # Converting html to lines lazily on render for easier debugging
     if (msg := app.selected_message) is not None:
-        msg_update_lines(msg, raw_mode=app.raw_mode)
+        msg.lines = msg_build_raw_lines(msg) if app.raw_mode else msg_build_lines(msg)
 
 
 def app_select_message(app: AppState, message: Optional[Message], show_pager: bool = False) -> None:
diff --git a/tests.py b/tests.py
@@ -16,7 +16,7 @@ class TestHtmlParser(unittest.TestCase):
         with open(html_path) as fp:
             html = fp.read()
 
-        actual = "\n".join(retronews.parse_html(html)).strip()
+        actual = "\n".join(retronews.parse_html(retronews.sanitize_text(html))).strip()
 
         if not os.path.exists(out_path):
             with open(out_path, "w") as fp: