ExternalVendorCode/lollms
1
0
Fork 0
mirror of https://github.com/ParisNeo/lollms.git synced 2024-12-24 06:46:40 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

enhanced security

Browse Source
This commit is contained in:
Saifeddine ALOUI 2024-07-06 03:48:01 +02:00
parent d3cfe3944f
commit 328b960a0d
1 changed files with 37 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

49
lollms/security.py
View File

@ -65,31 +65,56 @@ def sanitize_after_whitelisted_command(line, command):
# This means we should only return the part up to the whitelisted command # This means we should only return the part up to the whitelisted command
return line[:command_end_index + len(sanitized_rest)].strip() return line[:command_end_index + len(sanitized_rest)].strip()
if not(PackageManager.check_package_installed("defusedxml")):
PackageManager.install_or_update("defusedxml")
import defusedxml.ElementTree as ET
from defusedxml import ElementTree as ET
from io import StringIO
def sanitize_svg(svg_content): def sanitize_svg(svg_content):
try: try:
parser = ET.XMLParser(remove_comments=True, remove_pis=True) # Use defusedxml's parse function with a StringIO object
tree = ET.fromstring(svg_content, parser=parser) tree = ET.parse(StringIO(svg_content))
root = tree.getroot()
# Remove any script elements # Define a list of allowed elements
for script in tree.xpath('//svg:script', namespaces={'svg': 'http://www.w3.org/2000/svg'}): allowed_elements = {
parent = script.getparent() 'svg', 'g', 'path', 'circle', 'rect', 'line', 'polyline', 'polygon',
if parent is not None: 'text', 'tspan', 'defs', 'filter', 'feGaussianBlur', 'feMerge',
parent.remove(script) 'feMergeNode', 'linearGradient', 'radialGradient', 'stop'
}
# Remove any 'on*' event attributes # Define a list of allowed attributes
for element in tree.xpath('//*[@*[starts-with(name(), "on")]]'): allowed_attributes = {
'id', 'class', 'style', 'fill', 'stroke', 'stroke-width', 'cx', 'cy',
'r', 'x', 'y', 'width', 'height', 'd', 'transform', 'viewBox',
'xmlns', 'xmlns:xlink', 'version', 'stdDeviation', 'result', 'in',
'x1', 'y1', 'x2', 'y2', 'offset', 'stop-color', 'stop-opacity'
}
# Remove any disallowed elements
for element in root.iter():
if element.tag.split('}')[-1] not in allowed_elements:
parent = element.getparent()
if parent is not None:
parent.remove(element)
# Remove any disallowed attributes
for element in root.iter():
for attr in list(element.attrib): for attr in list(element.attrib):
if attr.startswith('on'): if attr not in allowed_attributes:
del element.attrib[attr] del element.attrib[attr]
# Convert the tree back to an SVG string # Convert the tree back to an SVG string
sanitized_svg = ET.tostring(tree, encoding='unicode', method='xml') sanitized_svg = ET.tostring(root, encoding='unicode', method='xml')
return sanitized_svg return sanitized_svg
except ET.XMLSyntaxError as e: except ET.ParseError as e:
raise ValueError("Invalid SVG content") from e raise ValueError("Invalid SVG content") from e
def sanitize_shell_code(code, whitelist=None): def sanitize_shell_code(code, whitelist=None):
""" """
Securely sanitizes a block of code by allowing commands from a provided whitelist, Securely sanitizes a block of code by allowing commands from a provided whitelist,