implement dompurify instead of bleach

modified:   app.py
	modified:   requirements.txt
	modified:   templates/chat.html

app.py

@@ -39,39 +39,6 @@ args = parser.parse_args()
 profile_name = args.profile
 
 
-def escape_except_code_and_pre(html_content):
-    soup = BeautifulSoup(html_content, "html.parser")
-
-    # Find all <code> and <pre> blocks and replace them with placeholders
-    placeholders = {}
-    for tag in soup.find_all(["code", "pre"]):
-        placeholder = f"PLACEHOLDER_{len(placeholders)}"
-        placeholders[placeholder] = str(tag)
-        tag.replace_with(placeholder)
-
-    # Convert the soup object back to a string
-    html_str = str(soup)
-
-    # Define bleach settings to allow iframes
-    tags = list(bleach.sanitizer.ALLOWED_TAGS) + ["iframe", "img"]
-    attributes = {
-        **bleach.sanitizer.ALLOWED_ATTRIBUTES,
-        "iframe": ["src", "width", "height", "frameborder", "allow", "allowfullscreen"],
-        "img": ["src", "alt", "width", "height", "title"],
-    }
-
-    # Sanitize the content using bleach
-    sanitized_content = bleach.clean(
-        html_str, tags=tags, attributes=attributes, strip=True
-    )
-
-    # Put back the <code> and <pre> blocks
-    for placeholder, original in placeholders.items():
-        sanitized_content = sanitized_content.replace(placeholder, original)
-
-    return sanitized_content
-
-
 class Message(db.Model):
     id = db.Column(db.Integer, primary_key=True)
     username = db.Column(db.String(128), nullable=False)
@@ -122,7 +89,7 @@ def on_join(data):
             {
                 "id": message.id,
                 "username": message.username,
-                "message": escape_except_code_and_pre(message.content),
+                "message": message.content,
             },
             room=request.sid,
         )
@@ -141,7 +108,7 @@ def handle_message(data):
     # Save the message to the database
     new_message = Message(
         username=data["username"],
-        content=escape_except_code_and_pre(data["message"]),
+        content=data["message"],
         room=data["room"],
     )
     db.session.add(new_message)
@@ -187,7 +154,7 @@ def handle_message(data):
                 data["username"],
                 data["room"],
                 data["message"],
-                model_name="gpt-4",
+                model_name="gpt-4-1106-preview",
             )
 
 
@@ -219,6 +186,7 @@ def chat_claude(username, room, message, model_name="anthropic.claude-v1"):
             "anthropic.claude-v1",
             "anthropic.claude-v2",
             "gpt-4",
+            "gpt-4-1106-preview",
         ]:
             chat_history += f"Assistant: {msg.username}: {msg.content}\n\n"
         else:
@@ -308,11 +276,15 @@ def chat_claude(username, room, message, model_name="anthropic.claude-v1"):
 
 
 def chat_gpt(username, room, message, model_name="gpt-3.5-turbo"):
+    limit = 15
+    if model_name == "gpt-4-1106-preview":
+        limit = 1000
+
     with app.app_context():
         last_messages = (
             Message.query.filter_by(room=room)
             .order_by(Message.id.desc())
-            .limit(10)
+            .limit(limit)
             .all()
         )
 
@@ -324,6 +296,7 @@ def chat_gpt(username, room, message, model_name="gpt-3.5-turbo"):
                 or msg.username == "anthropic.claude-v1"
                 or msg.username == "anthropic.claude-v2"
                 or msg.username == "gpt-4"
+                or msg.username == "gpt-4-1106-preview"
             )
             else "user",
             "content": f"{msg.username}: {msg.content}",

requirements.txt

@@ -7,6 +7,3 @@ openai
 Flask-SQLAlchemy
 
 boto3
-
-bleach
-beautifulsoup4

templates/chat.html

@@ -12,6 +12,8 @@
 
     <script src="https://cdnjs.cloudflare.com/ajax/libs/socket.io/4.0.1/socket.io.min.js"></script>
     <script src="https://cdnjs.cloudflare.com/ajax/libs/marked/9.1.2/marked.min.js" integrity="sha512-rfX4p3RNnxdwLT3wWP1K0NR3ztTobn+sISlT9WhxDDK00zNYbQ6MCHA5OHm0hqKAzEMXYCgFrp8iY/ER5MkXqA==" crossorigin="anonymous" referrerpolicy="no-referrer"></script>
+    <script src="https://cdn.jsdelivr.net/npm/dompurify@2/dist/purify.min.js"></script>
+
 
     <link rel="icon" href="{{ url_for('static', filename='favicon.ico') }}">
 
@@ -87,6 +89,15 @@ const urlParams = new URLSearchParams(window.location.search);
 const username = urlParams.get('username');
 const room = '{{ room }}';
 
+const dompurify_config = {
+  ADD_TAGS: ['iframe', 'img'],
+  FORBID_TAGS: ['form'], 
+  ALLOWED_ATTR: [
+    'src', 'width', 'height', 'frameborder', 'allowfullscreen', 
+    'alt', 'class', 'title', 'style'
+  ]
+};
+
 // Function to handle sending the message
 function sendMessage() {
     const message = document.getElementById('message').value;
@@ -126,7 +137,7 @@ socket.on('message', (data) => {
     messageWrapper.id = "message-" + data.id;
 
     const newMessage = document.createElement('p');
-    newMessage.innerHTML = marked.marked(data.content);
+    newMessage.innerHTML = DOMPurify.sanitize(marked.marked(data.content), dompurify_config);
 
     // Check if the message is NOT a join notification
     if (!data.content.includes('has joined the room.') && !data.content.includes('Processing...')) {
@@ -155,7 +166,7 @@ socket.on('previous_messages', (data) => {
     messageWrapper.id = "message-" + data.id;
 
     const newMessage = document.createElement('p');
-    newMessage.innerHTML = marked.marked(`**${data.username}:**\n\n${data.message}`);
+    newMessage.innerHTML = DOMPurify.sanitize(marked.marked(`**${data.username}:**\n\n${data.message}`), dompurify_config);
 
     // Create the delete button
     const deleteButton = document.createElement('button');
@@ -222,7 +233,7 @@ socket.on('message_chunk', (data) => {
     messageBuffers[data.id] += data.content;
 
     // Process the entire buffer with marked and set it as the content of the target element
-    targetMessageElement.innerHTML = marked.marked(messageBuffers[data.id]);
+    targetMessageElement.innerHTML = DOMPurify.sanitize(marked.marked(messageBuffers[data.id]), dompurify_config);
 
     // Apply syntax highlighting to code blocks within the content
     targetMessageElement.querySelectorAll('pre code').forEach((block) => {