feat(mcp): Protect from Prompt Injection by removing hidden content
Solves #214 as well
This commit is contained in:
+33
-1
@@ -2,7 +2,7 @@
|
|||||||
from sys import stderr
|
from sys import stderr
|
||||||
from copy import deepcopy
|
from copy import deepcopy
|
||||||
from functools import wraps
|
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 collections import namedtuple
|
||||||
from shlex import split as shlex_split
|
from shlex import split as shlex_split
|
||||||
from inspect import signature, Parameter
|
from inspect import signature, Parameter
|
||||||
@@ -21,6 +21,7 @@ from logging import (
|
|||||||
getLevelName,
|
getLevelName,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
from lxml.etree import XPath
|
||||||
from orjson import loads as json_loads, JSONDecodeError
|
from orjson import loads as json_loads, JSONDecodeError
|
||||||
|
|
||||||
from ._shell_signatures import Signatures_map
|
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
|
# Suppress exit on error to handle parsing errors gracefully
|
||||||
class NoExitArgumentParser(ArgumentParser): # pragma: no cover
|
class NoExitArgumentParser(ArgumentParser): # pragma: no cover
|
||||||
@@ -580,6 +594,23 @@ class Convertor:
|
|||||||
element.drop_tree()
|
element.drop_tree()
|
||||||
return Selector(root=clean_root, url=page.url)
|
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
|
@classmethod
|
||||||
def _extract_content(
|
def _extract_content(
|
||||||
cls,
|
cls,
|
||||||
@@ -597,6 +628,7 @@ class Convertor:
|
|||||||
if main_content_only:
|
if main_content_only:
|
||||||
page = cast(Selector, page.css("body").first) or page
|
page = cast(Selector, page.css("body").first) or page
|
||||||
page = cls._strip_noise_tags(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))
|
pages = [page] if not css_selector else cast(Selectors, page.css(css_selector))
|
||||||
for page in pages:
|
for page in pages:
|
||||||
|
|||||||
Reference in New Issue
Block a user