From f91aac1e11a3f4d39d615748a9afe64b6f54a7fe Mon Sep 17 00:00:00 2001
From: KnugiHK <24708955+KnugiHK@users.noreply.github.com>
Date: Sun, 11 Jul 2021 18:17:06 +0800
Subject: [PATCH] Implementing newline to <br>

---
 Whatsapp_Chat_Exporter/extract.py        | 20 ++++++++++++++++++--
 Whatsapp_Chat_Exporter/extract_iphone.py | 17 +++++++++++++++++
 Whatsapp_Chat_Exporter/whatsapp.html     |  4 ++--
 setup.py                                 |  3 ++-
 4 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/Whatsapp_Chat_Exporter/extract.py b/Whatsapp_Chat_Exporter/extract.py
index 5096f38..ce32001 100644
--- a/Whatsapp_Chat_Exporter/extract.py
+++ b/Whatsapp_Chat_Exporter/extract.py
@@ -8,6 +8,8 @@ import requests
 import shutil
 import re
 import pkgutil
+from bleach import clean as sanitize
+from markupsafe import Markup
 from datetime import datetime
 from mimetypes import MimeTypes
 try:
@@ -19,6 +21,10 @@ else:
     support_backup = True
 
 
+def sanitize_except(html):
+    return Markup(sanitize(html, tags=["br"]))
+
+
 def determine_day(last, current):
     last = datetime.fromtimestamp(last).date()
     current = datetime.fromtimestamp(current).date()
@@ -193,6 +199,11 @@ def messages(db, data):
                         msg = "{ Location shared: "f"{content[10], content[11]}"" }"
                     else:
                         msg = content[4]
+                        if msg is not None:
+                            if "\r\n" in msg:
+                                msg = msg.replace("\r\n", "<br>")
+                            if "\n" in msg:
+                                msg = msg.replace("\n", "<br>")
             else:
                 if content[5] == 0 and content[6] == 7:
                     msg = "{Message deleted}"
@@ -201,6 +212,11 @@ def messages(db, data):
                         msg = "{ Location shared: "f"{content[10], content[11]}"" }"
                     else:
                         msg = content[4]
+                        if msg is not None:
+                            if "\r\n" in msg:
+                                msg = msg.replace("\r\n", "<br>")
+                            if "\n" in msg:
+                                msg = msg.replace("\n", "<br>")
 
             data[content[0]]["messages"][content[1]]["data"] = msg
 
@@ -208,8 +224,7 @@ def messages(db, data):
         if i % 1000 == 0:
             print(f"Gathering messages...({i}/{total_row_number})", end="\r")
         content = c.fetchone()
-    print(
-        f"Gathering messages...({total_row_number}/{total_row_number})", end="\r")
+    print(f"Gathering messages...({total_row_number}/{total_row_number})", end="\r")
 
 
 def media(db, data, media_folder):
@@ -304,6 +319,7 @@ def create_html(data, output_folder, template=None):
     templateLoader = jinja2.FileSystemLoader(searchpath=template_dir)
     templateEnv = jinja2.Environment(loader=templateLoader)
     templateEnv.globals.update(determine_day=determine_day)
+    templateEnv.filters['sanitize_except'] = sanitize_except
     template = templateEnv.get_template(template_file)
 
     total_row_number = len(data)
diff --git a/Whatsapp_Chat_Exporter/extract_iphone.py b/Whatsapp_Chat_Exporter/extract_iphone.py
index 614abff..93bf672 100644
--- a/Whatsapp_Chat_Exporter/extract_iphone.py
+++ b/Whatsapp_Chat_Exporter/extract_iphone.py
@@ -7,12 +7,18 @@ import os
 import requests
 import shutil
 import pkgutil
+from bleach import clean as sanitize
+from markupsafe import Markup
 from datetime import datetime
 from mimetypes import MimeTypes
 
 APPLE_TIME = datetime.timestamp(datetime(2001, 1, 1))
 
 
+def sanitize_except(html):
+    return Markup(sanitize(html, tags=["br"]))
+
+
 def determine_day(last, current):
     last = datetime.fromtimestamp(last).date()
     current = datetime.fromtimestamp(current).date()
@@ -102,11 +108,21 @@ def messages(db, data):
                     msg = "{Message deleted}"
                 else:
                     msg = content[4]
+                    if msg is not None:
+                        if "\r\n" in msg:
+                            msg = msg.replace("\r\n", "<br>")
+                        if "\n" in msg:
+                            msg = msg.replace("\n", "<br>")
             else:
                 if content[5] == 14:
                     msg = "{Message deleted}"
                 else:
                     msg = content[4]
+                    if msg is not None:
+                        if "\r\n" in msg:
+                            msg = msg.replace("\r\n", "<br>")
+                        if "\n" in msg:
+                            msg = msg.replace("\n", "<br>")
             data[content[0]]["messages"][content[1]]["data"] = msg
         i += 1
         if i % 1000 == 0:
@@ -216,6 +232,7 @@ def create_html(data, output_folder, template=None):
     templateLoader = jinja2.FileSystemLoader(searchpath=template_dir)
     templateEnv = jinja2.Environment(loader=templateLoader)
     templateEnv.globals.update(determine_day=determine_day)
+    templateEnv.filters['sanitize_except'] = sanitize_except
     template = templateEnv.get_template(template_file)
 
     total_row_number = len(data)
diff --git a/Whatsapp_Chat_Exporter/whatsapp.html b/Whatsapp_Chat_Exporter/whatsapp.html
index 7bea23b..8273da6 100644
--- a/Whatsapp_Chat_Exporter/whatsapp.html
+++ b/Whatsapp_Chat_Exporter/whatsapp.html
@@ -73,7 +73,7 @@
 								</div>
 							{% endif %}
 							{% if msg.media == false %}
-								{% filter escape %}{{ msg.data or "{This message is not supported yet}" | replace('\n', '<br>') }}{% endfilter %}
+								{{ msg.data or "{This message is not supported yet}" | sanitize_except() }}
 							{% else %}
 								{% if "image/" in msg.mime %}
 								<a href="{{ msg.data }}"><img src="{{ msg.data }}" /></a>
@@ -121,7 +121,7 @@
 								</div>
 							{% endif %}
 							{% if msg.media == false %}
-								{% filter escape %}{{ msg.data or "{This message is not supported yet}" }}{% endfilter %}
+							{{ msg.data or "{This message is not supported yet}" | sanitize_except() }}
 							{% else %}
 								{% if "image/" in msg.mime %}
 								<a href="{{ msg.data }}"><img src="{{ msg.data }}" /></a>
diff --git a/setup.py b/setup.py
index af396c7..6633374 100644
--- a/setup.py
+++ b/setup.py
@@ -37,7 +37,8 @@ setuptools.setup(
     ],
     python_requires='>=3.7',
     install_requires=[
-       'jinja2'
+       'jinja2',
+       'bleach'
     ],
     extras_require={
         'android_backup':  ["pycryptodome"]