From 251d90301c2bbaaf2c46362986a321fc70712ecc Mon Sep 17 00:00:00 2001
From: Tharo <17233964+Thar0@users.noreply.github.com>
Date: Mon, 14 Feb 2022 23:26:03 +0000
Subject: [PATCH] Make text headers encoding more robust (#1145)

* Improve encoding stage of text headers

* Fix r string

* Remove another unnecessary backslash
---
 Makefile        |  3 ++-
 tools/msgenc.py | 46 ++++++++++++++++++++++++++++------------------
 2 files changed, 30 insertions(+), 19 deletions(-)

diff --git a/Makefile b/Makefile
index 1fcd8a66e5..86285c9659 100644
--- a/Makefile
+++ b/Makefile
@@ -246,7 +246,8 @@ build/data/%.o: data/%.s
 	$(AS) $(ASFLAGS) $< -o $@
 
 build/assets/text/%.enc.h: assets/text/%.h assets/text/charmap.txt
-	python3 tools/msgenc.py assets/text/charmap.txt $< $@
+	$(CPP) -P -dD -fpreprocessed $< > $(@:.enc.h=.h)
+	python3 tools/msgenc.py assets/text/charmap.txt $(@:.enc.h=.h) $@
 
 build/assets/text/fra_message_data_static.o: build/assets/text/message_data.enc.h
 build/assets/text/ger_message_data_static.o: build/assets/text/message_data.enc.h
diff --git a/tools/msgenc.py b/tools/msgenc.py
index 80cdd653c0..b2b9cab168 100644
--- a/tools/msgenc.py
+++ b/tools/msgenc.py
@@ -5,41 +5,51 @@
 
 import argparse, ast, re
 
-charmap = {}
-
-string_regex = re.compile(r"([\"'`])(?:[\s\S])*?(?:(?<!\\)\1)")
-
 def read_charmap(path):
-    global charmap
-
     with open(path) as infile:
         charmap = infile.read()
 
     charmap = ast.literal_eval(charmap)
-    charmap = dict([(repr(k)[1:-1],chr(v)) for k,v in charmap.items()])
+    charmap = { repr(k)[1:-1] : chr(v) for k,v in charmap.items() }
 
-def cvt_str(m):
-    string = m.group(0)
+    return charmap
 
-    for orig,char in charmap.items():
-        string = string.replace(orig, char)
+def convert_text(text, charmap):
+    def cvt_str(m):
+        string = m.group(0)
 
-    return string
+        for orig,char in charmap.items():
+            string = string.replace(orig, char)
 
-if __name__ == "__main__":
+        return string
+
+    # Naive string matcher, assumes single line strings and no comments, handles escaped quotations
+    string_regex = re.compile(r'"((?:[^\\"\n]|\\.)*)"')
+
+    # Collapse escaped newlines
+    text = text.replace("\\\n", "")
+    # Encode according to charmap
+    text = re.sub(string_regex, cvt_str, text)
+
+    return text
+
+def main():
     parser = argparse.ArgumentParser(description="Encode message_data_static text headers")
     parser.add_argument("charmap", help="path to charmap file specifying custom encoding elements")
     parser.add_argument("input", help="path to file to be encoded")
     parser.add_argument("output", help="encoded file")
     args = parser.parse_args()
 
-    read_charmap(args.charmap)
+    charmap = read_charmap(args.charmap)
 
-    contents = ""
+    text = ""
     with open(args.input, "r") as infile:
-        contents = infile.read()
+        text = infile.read()
 
-    contents = re.sub(string_regex, cvt_str, contents)
+    text = convert_text(text, charmap)
 
     with open(args.output, "w", encoding="raw_unicode_escape") as outfile:
-        outfile.write(contents)
+        outfile.write(text)
+
+if __name__ == "__main__":
+    main()