From e217036f3a6c8b45aa838ef48c6f649ac7c73163 Mon Sep 17 00:00:00 2001
From: SimoneMariaRomeo <180769497+SimoneMariaRomeo@users.noreply.github.com>
Date: Wed, 13 May 2026 02:15:03 +0700
Subject: [PATCH] Escape profile badge custom labels

---
 profile_badge_generator.py                    | 54 ++++++++++++++-----
 .../test_profile_badge_generator_security.py  | 52 ++++++++++++++++++
 2 files changed, 94 insertions(+), 12 deletions(-)
 create mode 100644 tests/test_profile_badge_generator_security.py

diff --git a/profile_badge_generator.py b/profile_badge_generator.py
index 529925bde..59e08b70d 100644
--- a/profile_badge_generator.py
+++ b/profile_badge_generator.py
@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: MIT
 
 from flask import Flask, request, jsonify, render_template_string
+import html as html_utils
 import sqlite3
 import json
 import urllib.parse
@@ -109,7 +110,12 @@ def badge_generator():
             .then(response => response.json())
             .then(data => {
                 if (data.success) {
-                    document.getElementById('badgePreview').innerHTML = data.preview_html;
+                    const badgePreview = document.getElementById('badgePreview');
+                    badgePreview.replaceChildren();
+                    const previewImage = document.createElement('img');
+                    previewImage.src = data.shield_url;
+                    previewImage.alt = data.alt_text || 'RustChain Badge';
+                    badgePreview.appendChild(previewImage);
                     document.getElementById('markdownCode').textContent = data.markdown;
                     document.getElementById('htmlCode').textContent = data.html;
                     document.getElementById('result').style.display = 'block';
@@ -123,15 +129,33 @@ def badge_generator():
     '''
     return render_template_string(html)
 
+
+def text_field(data, name, default=''):
+    value = data.get(name, default)
+    if value is None:
+        return ''
+    return str(value).strip()
+
+
+def escape_markdown_alt(text):
+    return (
+        text.replace('\\', '\\\\')
+        .replace('[', '\\[')
+        .replace(']', '\\]')
+        .replace('\n', ' ')
+        .replace('\r', ' ')
+    )
+
 @app.route('/api/badge/create', methods=['POST'])
 def create_badge():
     init_badge_db()
-    data = request.get_json()
+    raw_data = request.get_json(silent=True) or {}
+    data = raw_data if isinstance(raw_data, dict) else {}
     
-    username = data.get('username', '').strip()
-    wallet = data.get('wallet', '').strip()
-    badge_type = data.get('badge_type', 'contributor')
-    custom_message = data.get('custom_message', '').strip()
+    username = text_field(data, 'username')
+    wallet = text_field(data, 'wallet')
+    badge_type = text_field(data, 'badge_type', 'contributor')
+    custom_message = text_field(data, 'custom_message')
     
     if not username:
         return jsonify({'success': False, 'error': 'Username required'})
@@ -146,12 +170,17 @@ def create_badge():
     color = badge_colors.get(badge_type, 'blue')
     label = custom_message if custom_message else badge_type.replace('-', ' ').title()
     
-    shield_url = f"https://img.shields.io/badge/RustChain-{urllib.parse.quote(label)}-{color}"
+    shield_url = f"https://img.shields.io/badge/RustChain-{urllib.parse.quote(label, safe='')}-{color}"
     repo_url = "https://github.com/Scottcjn/Rustchain"
+    alt_text = f"RustChain {label}"
+    html_alt_text = html_utils.escape(alt_text, quote=True)
+    html_shield_url = html_utils.escape(shield_url, quote=True)
+    html_repo_url = html_utils.escape(repo_url, quote=True)
+    markdown_alt_text = escape_markdown_alt(alt_text)
     
-    markdown = f"[![RustChain {label}]({shield_url})]({repo_url})"
-    html = f'<a href="{repo_url}"><img src="{shield_url}" alt="RustChain {label}"></a>'
-    preview_html = f'<img src="{shield_url}" alt="RustChain {label}">'
+    markdown = f"[![{markdown_alt_text}]({shield_url})]({repo_url})"
+    html = f'<a href="{html_repo_url}"><img src="{html_shield_url}" alt="{html_alt_text}"></a>'
+    preview_html = f'<img src="{html_shield_url}" alt="{html_alt_text}">'
     
     with sqlite3.connect(DB_PATH) as conn:
         cursor = conn.cursor()
@@ -167,7 +196,8 @@ def create_badge():
         'markdown': markdown,
         'html': html,
         'preview_html': preview_html,
-        'shield_url': shield_url
+        'shield_url': shield_url,
+        'alt_text': alt_text
     })
 
 @app.route('/api/badge/stats')
@@ -216,4 +246,4 @@ def list_badges():
 
 if __name__ == '__main__':
     init_badge_db()
-    app.run(debug=True, port=5003)
\ No newline at end of file
+    app.run(debug=True, port=5003)
diff --git a/tests/test_profile_badge_generator_security.py b/tests/test_profile_badge_generator_security.py
new file mode 100644
index 000000000..1c1710adc
--- /dev/null
+++ b/tests/test_profile_badge_generator_security.py
@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: MIT
+from pathlib import Path
+import importlib.util
+
+
+MODULE_PATH = Path(__file__).resolve().parents[1] / "profile_badge_generator.py"
+
+
+def load_profile_badge_module(tmp_path):
+    spec = importlib.util.spec_from_file_location("profile_badge_generator_under_test", MODULE_PATH)
+    module = importlib.util.module_from_spec(spec)
+    assert spec.loader is not None
+    spec.loader.exec_module(module)
+    module.DB_PATH = str(tmp_path / "profile_badges.db")
+    return module
+
+
+def test_custom_message_is_escaped_in_generated_badge_markup(tmp_path):
+    module = load_profile_badge_module(tmp_path)
+    client = module.app.test_client()
+    payload = 'Active"] <script>alert(1)</script> / badge'
+
+    response = client.post(
+        "/api/badge/create",
+        json={
+            "username": "alice",
+            "badge_type": "contributor",
+            "custom_message": payload,
+        },
+    )
+
+    assert response.status_code == 200
+    data = response.get_json()
+    assert data["success"] is True
+    assert "%22%5D%20%3Cscript%3Ealert%281%29%3C%2Fscript%3E%20%2F%20badge" in data["shield_url"]
+    assert "<script>" not in data["html"]
+    assert "&lt;script&gt;alert(1)&lt;/script&gt;" in data["html"]
+    assert 'Active&quot;] &lt;script&gt;' in data["preview_html"]
+    assert 'Active"] <script>' not in data["preview_html"]
+    assert 'RustChain Active"\\]' in data["markdown"]
+    assert 'RustChain Active"]' not in data["markdown"]
+
+
+def test_badge_generator_preview_uses_dom_api_not_returned_html():
+    source = MODULE_PATH.read_text(encoding="utf-8")
+
+    assert "badgePreview.replaceChildren();" in source
+    assert "document.createElement('img')" in source
+    assert "previewImage.src = data.shield_url;" in source
+    assert "previewImage.alt = data.alt_text || 'RustChain Badge';" in source
+    assert "badgePreview.appendChild(previewImage);" in source
+    assert "badgePreview').innerHTML = data.preview_html" not in source