From 9d532c2a76c69b634a698a7f482c1cd90b9829bb Mon Sep 17 00:00:00 2001 From: Saifeddine ALOUI Date: Thu, 28 Mar 2024 23:59:15 +0100 Subject: [PATCH] enhanced security --- lollms/internet.py | 15 ++++++++------- lollms/security.py | 6 +++--- 2 files changed, 11 insertions(+), 10 deletions(-) diff --git a/lollms/internet.py b/lollms/internet.py index e3262ce..0b564bc 100644 --- a/lollms/internet.py +++ b/lollms/internet.py @@ -105,7 +105,7 @@ def press_buttons(driver, buttons_to_press=['accept']): except: ASCIIColors.warning(f"Couldn't press button {button_to_press} in this page.") -def scrape_and_save(url, file_path=None, lollms_com=None, chromedriver_path=None, wait_step_delay=1, buttons_to_press=['accept']): +def scrape_and_save(url, file_path=None, lollms_com=None, chromedriver_path=None, wait_step_delay=1, buttons_to_press=['accept'], max_size=None): if not PackageManager.check_package_installed("selenium"): PackageManager.install_package("selenium") if not PackageManager.check_package_installed("bs4"): @@ -135,12 +135,13 @@ def scrape_and_save(url, file_path=None, lollms_com=None, chromedriver_path=None if file_path: - # Save the text content as a text file - with open(file_path, 'w', encoding="utf-8") as file: - file.write(text_content) - if lollms_com: - lollms_com.info(f"Webpage content saved to {file_path}") - + if max_size and text_content< max_size: + # Save the text content as a text file + with open(file_path, 'w', encoding="utf-8") as file: + file.write(text_content) + if lollms_com: + lollms_com.info(f"Webpage content saved to {file_path}") + # Close the driver driver.quit() diff --git a/lollms/security.py b/lollms/security.py index c363b07..8a482ec 100644 --- a/lollms/security.py +++ b/lollms/security.py @@ -46,9 +46,9 @@ def sanitize_path_from_endpoint(path: str, error_text="A suspected LFI attack de return path -def forbid_remote_access(lollmsElfServer): - if lollmsElfServer.config.host!="localhost" and lollmsElfServer.config.host!="127.0.0.1": - raise Exception("This functionality is forbidden if the server is exposed") +def forbid_remote_access(lollmsElfServer, exception_text = "This functionality is forbidden if the server is exposed"): + if not lollmsElfServer.config.force_accept_remote_access and lollmsElfServer.config.host!="localhost" and lollmsElfServer.config.host!="127.0.0.1": + raise Exception(exception_text) def validate_path(path, allowed_paths:List[str|Path]): # Convert the path to an absolute path