Skip to content

Commit

Permalink
enhanced security
Browse files Browse the repository at this point in the history
  • Loading branch information
ParisNeo committed Jul 6, 2024
1 parent d3cfe39 commit 328b960
Showing 1 changed file with 39 additions and 14 deletions.
53 changes: 39 additions & 14 deletions lollms/security.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,31 +65,56 @@ def sanitize_after_whitelisted_command(line, command):
# This means we should only return the part up to the whitelisted command
return line[:command_end_index + len(sanitized_rest)].strip()

if not(PackageManager.check_package_installed("defusedxml")):
PackageManager.install_or_update("defusedxml")

def sanitize_svg(svg_content):
try:
parser = ET.XMLParser(remove_comments=True, remove_pis=True)
tree = ET.fromstring(svg_content, parser=parser)
import defusedxml.ElementTree as ET

# Remove any script elements
for script in tree.xpath('//svg:script', namespaces={'svg': 'http://www.w3.org/2000/svg'}):
parent = script.getparent()
if parent is not None:
parent.remove(script)
from defusedxml import ElementTree as ET
from io import StringIO

# Remove any 'on*' event attributes
for element in tree.xpath('//*[@*[starts-with(name(), "on")]]'):
def sanitize_svg(svg_content):
try:
# Use defusedxml's parse function with a StringIO object
tree = ET.parse(StringIO(svg_content))
root = tree.getroot()

# Define a list of allowed elements
allowed_elements = {
'svg', 'g', 'path', 'circle', 'rect', 'line', 'polyline', 'polygon',
'text', 'tspan', 'defs', 'filter', 'feGaussianBlur', 'feMerge',
'feMergeNode', 'linearGradient', 'radialGradient', 'stop'
}

# Define a list of allowed attributes
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):
if attr.startswith('on'):
if attr not in allowed_attributes:
del element.attrib[attr]

# 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
except ET.XMLSyntaxError as e:
except ET.ParseError as e:
raise ValueError("Invalid SVG content") from e



def sanitize_shell_code(code, whitelist=None):
"""
Securely sanitizes a block of code by allowing commands from a provided whitelist,
Expand Down

0 comments on commit 328b960

Please sign in to comment.