diff --git a/app.py b/app.py index 18a99e6..56e9749 100644 --- a/app.py +++ b/app.py @@ -14,6 +14,9 @@ import time import boto3 import json +from bs4 import BeautifulSoup +import bleach + app = Flask(__name__) app.config["SECRET_KEY"] = "your_secret_key" @@ -36,6 +39,38 @@ args = parser.parse_args() profile_name = args.profile +def escape_except_code_and_pre(html_content): + soup = BeautifulSoup(html_content, "html.parser") + + # Find all and
 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"]
+    attributes = {
+        **bleach.sanitizer.ALLOWED_ATTRIBUTES,
+        "iframe": ["src", "width", "height", "frameborder", "allow", "allowfullscreen"],
+    }
+
+    # Sanitize the content using bleach
+    sanitized_content = bleach.clean(
+        html_str, tags=tags, attributes=attributes, strip=True
+    )
+
+    # Put back the  and 
 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)
@@ -86,7 +121,7 @@ def on_join(data):
             {
                 "id": message.id,
                 "username": message.username,
-                "message": message.content,
+                "message": escape_except_code_and_pre(message.content),
             },
             room=request.sid,
         )
@@ -104,7 +139,9 @@ def on_join(data):
 def handle_message(data):
     # Save the message to the database
     new_message = Message(
-        username=data["username"], content=data["message"], room=data["room"]
+        username=data["username"],
+        content=escape_except_code_and_pre(data["message"]),
+        room=data["room"],
     )
     db.session.add(new_message)
     db.session.commit()
diff --git a/requirements.txt b/requirements.txt
index dd99b4c..7ccceb3 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,3 +7,6 @@ openai
 Flask-SQLAlchemy
 
 boto3
+
+bleach
+beautifulsoup4