Make text headers encoding more robust (#1145)

* Improve encoding stage of text headers

* Fix r string

* Remove another unnecessary backslash

1 files changed, 28 insertions, 18 deletions

diff --git a/tools/msgenc.py b/tools/msgenc.py
index 80cdd653c..b2b9cab16 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()