feat(mcp): Protect from Prompt Injection by removing hidden content

Solves #214 as well
This commit is contained in:
Karim shoair
2026-03-30 02:31:14 +02:00
parent bcd39d57c0
commit 375951bd49
+33 -1
View File
@@ -2,7 +2,7 @@
from sys import stderr
from copy import deepcopy
from functools import wraps
from re import sub as re_sub
from re import sub as re_sub, compile as re_compile
from collections import namedtuple
from shlex import split as shlex_split
from inspect import signature, Parameter
@@ -21,6 +21,7 @@ from logging import (
getLevelName,
)
from lxml.etree import XPath
from orjson import loads as json_loads, JSONDecodeError
from ._shell_signatures import Signatures_map
@@ -67,6 +68,19 @@ Request = namedtuple(
],
)
# Precompiled for the prompt injection sanitizer
_HIDDEN_XPATH = XPath(
'.//*[contains(@style,"display:none") or contains(@style,"display: none")'
' or contains(@style,"visibility:hidden") or contains(@style,"visibility: hidden")'
' or contains(@style,"opacity:0") or contains(@style,"opacity: 0")'
' or contains(@style,"font-size:0") or contains(@style,"font-size: 0")'
' or contains(@style,"height:0") or contains(@style,"height: 0")'
' or contains(@style,"width:0") or contains(@style,"width: 0")]'
" | .//*[@aria-hidden='true']"
" | .//template"
)
_ZWC_PATTERN = re_compile(r"[\u200b\u200c\u200d\ufeff\u2060\u180e]")
# Suppress exit on error to handle parsing errors gracefully
class NoExitArgumentParser(ArgumentParser): # pragma: no cover
@@ -580,6 +594,23 @@ class Convertor:
element.drop_tree()
return Selector(root=clean_root, url=page.url)
@classmethod
def _sanitize_for_ai(cls, page: Selector) -> Selector:
"""Strip hidden content that could be used for prompt injection.
Removes CSS-hidden elements, aria-hidden elements, <template> tags,
HTML comments, and zero-width Unicode characters.
"""
clean_root = deepcopy(page._root)
for element in cast(list, _HIDDEN_XPATH(clean_root)):
element.drop_tree()
for element in clean_root.iter():
if element.text:
element.text = _ZWC_PATTERN.sub("", element.text)
if element.tail:
element.tail = _ZWC_PATTERN.sub("", element.tail)
return Selector(root=clean_root, url=page.url, keep_comments=False)
@classmethod
def _extract_content(
cls,
@@ -597,6 +628,7 @@ class Convertor:
if main_content_only:
page = cast(Selector, page.css("body").first) or page
page = cls._strip_noise_tags(page)
page = cls._sanitize_for_ai(page)
pages = [page] if not css_selector else cast(Selectors, page.css(css_selector))
for page in pages: