import html
import bleach # pip install bleach

def sanitize_for_display(user_input_html):
    """
    Sanitizes user-provided HTML content to prevent XSS attacks
    before displaying it on a web page.
    """
    # Recommended tags and attributes that are generally safe.
    # Customize this list based on your application's needs.
    allowed_tags = [
        'a', 'abbr', 'acronym', 'b', 'blockquote', 'code',
        'em', 'i', 'li', 'ol', 'p', 'strong', 'ul', 'br',
        'h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'div', 'span'
    ]
    allowed_attrs = {
        'a': ['href', 'title'],
        'abbr': ['title'],
        'acronym': ['title'],
    }

    # Clean the HTML using bleach.
    # strip=True removes disallowed tags and their content.
    # If strip=False, disallowed tags are removed, but their content is kept.
    clean_html = bleach.clean(
        user_input_html,
        tags=allowed_tags,
        attributes=allowed_attrs,
        strip=True # Or strip=False if you want to keep content inside stripped tags
    )
    return clean_html

def escape_plain_text(user_input_text):
    """
    Escapes special characters in plain text to prevent them from being
    interpreted as HTML. Use this for displaying text that should *never*
    contain HTML.
    """
    return html.escape(user_input_text)

# --- Example Usage ---

# Scenario 1: User inputs potentially malicious HTML
malicious_html_input = (
    "<script>alert('XSS Attack!');</script>"
    "<p>Hello, <img src=x onerror=alert('Image XSS')>world!</p>"
    "<a href=\"javascript:alert('Link XSS')\">Click me</a>"
    "<h1>My Blog Post</h1>"
)
sanitized_output = sanitize_for_display(malicious_html_input)
print("--- Sanitized HTML Output (displaying allowed tags, removing malicious) ---")
print(sanitized_output)
# Expected: "Hello, world!<a href=\"javascript:alert('Link XSS')\">Click me</a><h1>My Blog Post</h1>"
# Note: bleach by default removes href="javascript:..."
# The output will be: <p>Hello, world!</p><a>Click me</a><h1>My Blog Post</h1>

print("
--- Corrected expected output for bleach ---")
malicious_html_input_bleach_test = (
    "<script>alert('XSS Attack!');</script>"
    "<p>Hello, <img src=x onerror=alert('Image XSS')>world!</p>"
    "<a href=\"javascript:alert('Link XSS')\">Click me</a>"
    "<h1>My Blog Post</h1>"
)
# The default behavior of bleach for a tag with href="javascript:..." is to remove the href attribute.
# It will become <a>Click me</a>. If the tag itself is not allowed, it's stripped.
# The img tag is not in allowed_tags, so it will be stripped.
# Script tag is not in allowed_tags, so it will be stripped.
print(sanitize_for_display(malicious_html_input_bleach_test))


# Scenario 2: User inputs plain text that might contain HTML-like characters
plain_text_input = "User's comment: <p>This is great!</p> & that's all."
escaped_text_output = escape_plain_text(plain_text_input)
print("
--- Escaped Plain Text Output (for non-HTML content) ---")
print(escaped_text_output)
# Expected: "User&#x27;s comment: &lt;p&gt;This is great!&lt;/p&gt; &amp; that&#x27;s all."