From eda3af5f5c4ea9b2f3569f72f8d05989e29367fc Mon Sep 17 00:00:00 2001 From: Saifeddine ALOUI Date: Sun, 19 May 2024 17:38:38 +0200 Subject: [PATCH] Added sanitize_svg function and forbid_remote_access is used for setting personality settings --- lollms/security.py | 30 +++++++++++++++++++ .../endpoints/lollms_personalities_infos.py | 18 +++++++++-- lollms/utilities.py | 5 +++- 3 files changed, 50 insertions(+), 3 deletions(-) diff --git a/lollms/security.py b/lollms/security.py index 10a932c..f5f3fa4 100644 --- a/lollms/security.py +++ b/lollms/security.py @@ -8,6 +8,12 @@ import os import re import platform import string +from lollms.utilities import PackageManager + +if not PackageManager.check_package_installed("lxml"): + PackageManager.install_package("lxml") + +import lxml.etree as ET def check_access(lollmsElfServer, client_id): client = lollmsElfServer.session.get_client(client_id) @@ -60,6 +66,30 @@ def sanitize_after_whitelisted_command(line, command): return line[:command_end_index + len(sanitized_rest)].strip() +def sanitize_svg(svg_content): + try: + parser = ET.XMLParser(remove_comments=True, remove_pis=True) + tree = ET.fromstring(svg_content, parser=parser) + + # 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) + + # Remove any 'on*' event attributes + for element in tree.xpath('//*[@*[starts-with(name(), "on")]]'): + for attr in list(element.attrib): + if attr.startswith('on'): + del element.attrib[attr] + + # Convert the tree back to an SVG string + sanitized_svg = ET.tostring(tree, encoding='unicode', method='xml') + return sanitized_svg + except ET.XMLSyntaxError 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, diff --git a/lollms/server/endpoints/lollms_personalities_infos.py b/lollms/server/endpoints/lollms_personalities_infos.py index 0f28404..d192ebd 100644 --- a/lollms/server/endpoints/lollms_personalities_infos.py +++ b/lollms/server/endpoints/lollms_personalities_infos.py @@ -16,7 +16,7 @@ from lollms.server.elf_server import LOLLMSElfServer from lollms.personality import AIPersonality, InstallOption from ascii_colors import ASCIIColors from lollms.utilities import load_config, trace_exception, gc, show_yes_no_dialog -from lollms.security import check_access +from lollms.security import check_access, forbid_remote_access from pathlib import Path from typing import List, Optional import psutil @@ -329,6 +329,8 @@ def get_personality_config(data:PersonalityDataRequest): print("- Recovering personality config") category = sanitize_path(data.category) name = sanitize_path(data.name) + if category=="": + return {"status":False, "error":"category must not be empty."} package_path = f"{category}/{name}" if category=="custom_personalities": @@ -352,12 +354,15 @@ class PersonalityConfig(BaseModel): @router.post("/set_personality_config") def set_personality_config(data:PersonalityConfig): + forbid_remote_access(lollmsElfServer) check_access(lollmsElfServer, data.client_id) print("- Recovering personality config") category = sanitize_path(data.category) name = sanitize_path(data.name) config = data.config - + if category=="": + return {"status":False, "error":"category must not be empty."} + package_path = f"{category}/{name}" if category=="custom_personalities": package_full_path = lollmsElfServer.lollms_paths.custom_personalities_path/f"{name}" @@ -388,6 +393,8 @@ def mount_personality(data:PersonalityMountingInfos): category = sanitize_path(data.category) name = sanitize_path(data.folder) language = data.language #.get('language', None) + if category=="": + return {"status":False, "error":"category must not be empty."} package_path = f"{category}/{name}" if category=="custom_personalities": @@ -443,6 +450,8 @@ def remount_personality(data:PersonalityMountingInfos): name = sanitize_path(data.folder) language = data.language #.get('language', None) + if category=="": + return {"status":False, "error":"category must not be empty."} package_path = f"{category}/{name}" if category=="custom_personalities": @@ -497,6 +506,9 @@ def unmount_personality(data:PersonalityMountingInfos): name = sanitize_path(data.folder) language = data.language #.get('language', None) + if category=="": + return {"status":False, "error":"category must not be empty."} + try: personality_id = f"{category}/{name}" if language is None or language=="" else f"{category}/{name}:{language}" index = lollmsElfServer.config["personalities"].index(personality_id) @@ -595,6 +607,8 @@ def get_personality_settings(data:PersonalityMountingInfos): print("- Retreiving personality settings") category = sanitize_path(data.category) name = sanitize_path(data.folder) + if category=="": + return {"status":False, "error":"category must not be empty."} if category == "custom_personalities": personality_folder = lollmsElfServer.lollms_paths.personal_personalities_path/f"{name}" diff --git a/lollms/utilities.py b/lollms/utilities.py index 1e53004..f10fd2a 100644 --- a/lollms/utilities.py +++ b/lollms/utilities.py @@ -37,7 +37,6 @@ import git import mimetypes import subprocess -from lollms.security import sanitize_shell_code from functools import partial @@ -59,6 +58,7 @@ def install_conda_utility(): def create_conda_env(env_name, python_version): + from lollms.security import sanitize_shell_code env_name = sanitize_shell_code(env_name) python_version = sanitize_shell_code(python_version) # Activate the Conda environment @@ -132,6 +132,7 @@ def get_conda_path(): return Path(sys.executable).parent.parent.parent / "miniconda3" / "bin" / "conda" def environment_exists(env_name): + from lollms.security import sanitize_shell_code env_name = sanitize_shell_code(env_name) conda_path = get_conda_path() result = subprocess.run(f'{conda_path} env list --json', shell=True, capture_output=True, text=True) @@ -140,6 +141,7 @@ def environment_exists(env_name): return env_name in env_names def get_python_version(env_name): + from lollms.security import sanitize_shell_code env_name = sanitize_shell_code(env_name) conda_path = get_conda_path() if environment_exists(env_name): @@ -149,6 +151,7 @@ def get_python_version(env_name): return "Environment does not exist." def remove_environment(env_name): + from lollms.security import sanitize_shell_code env_name = sanitize_shell_code(env_name) conda_path = get_conda_path() if environment_exists(env_name):