From 104dc4e004fe532211069e0fc36b3e361046ebd8 Mon Sep 17 00:00:00 2001 From: Saifeddine ALOUI Date: Mon, 18 Mar 2024 22:30:13 +0100 Subject: [PATCH] upgraded --- lollms/security.py | 5 ++++ .../endpoints/lollms_configuration_infos.py | 5 +++- .../server/endpoints/lollms_skills_library.py | 24 +++++++++++++++---- lollms/server/endpoints/lollms_user.py | 5 +++- 4 files changed, 32 insertions(+), 7 deletions(-) diff --git a/lollms/security.py b/lollms/security.py index 8548e99..c363b07 100644 --- a/lollms/security.py +++ b/lollms/security.py @@ -7,6 +7,11 @@ from typing import List import os import re +def check_access(lollmsElfServer, client_id): + client = lollmsElfServer.session.get_client(client_id) + if not client: + raise HTTPException(status_code=400, detail=f"Not accessible without id") + return client def sanitize_path(path:str, allow_absolute_path:bool=False, error_text="Absolute database path detected", exception_text="Detected an attempt of path traversal. Are you kidding me?"): if path is None: diff --git a/lollms/server/endpoints/lollms_configuration_infos.py b/lollms/server/endpoints/lollms_configuration_infos.py index f22d803..825397e 100644 --- a/lollms/server/endpoints/lollms_configuration_infos.py +++ b/lollms/server/endpoints/lollms_configuration_infos.py @@ -15,6 +15,7 @@ from lollms.server.elf_server import LOLLMSElfServer from lollms.binding import BindingBuilder, InstallOption from ascii_colors import ASCIIColors from lollms.utilities import load_config, trace_exception, gc +from lollms.security import check_access from pathlib import Path from typing import List import json @@ -52,9 +53,9 @@ async def update_setting(request: Request): """ # Prevent all outsiders from sending something to this endpoint forbid_remote_access(lollmsElfServer) - try: config_data = (await request.json()) + check_access(lollmsElfServer, config_data["client_id"]) if "config" in config_data.keys(): config_data = config_data["config"] setting_name = config_data["setting_name"] @@ -122,6 +123,8 @@ async def update_setting(request: Request): lollmsElfServer.config.save_config() # Tell that the setting was changed return {'setting_name': setting_name, "status":True} + except HTTPException as ex: + raise ex except Exception as ex: trace_exception(ex) lollmsElfServer.error(ex) diff --git a/lollms/server/endpoints/lollms_skills_library.py b/lollms/server/endpoints/lollms_skills_library.py index 9d5193b..5d7e164 100644 --- a/lollms/server/endpoints/lollms_skills_library.py +++ b/lollms/server/endpoints/lollms_skills_library.py @@ -15,6 +15,7 @@ from lollms.utilities import detect_antiprompt, remove_text_from_string, trace_e from lollms.security import sanitize_path from ascii_colors import ASCIIColors from lollms.databases.discussions_database import DiscussionsDB, Discussion +from lollms.security import check_access from typing import List from safe_store.text_vectorizer import TextVectorizer, VectorizationMethod, VisualizationMethod @@ -25,7 +26,7 @@ from pathlib import Path router = APIRouter() lollmsElfServer:LOLLMSWebUI = LOLLMSWebUI.get_instance() -class clientInfos(BaseModel): +class ClientInfos(BaseModel): client_id: str class SkillInfos(BaseModel): @@ -39,42 +40,55 @@ class SkillUpdateInfos(BaseModel): title: str content: str +class DeleteSkillInfos(BaseModel): + client_id: str + skill_id: int + class CategoryData(BaseModel): client_id: str category: str @router.post("/get_skills_library") -def get_skills_library_categories(discussionInfos:clientInfos): +def get_skills_library_categories(discussionInfos:ClientInfos): return {"status":True, "entries":lollmsElfServer.skills_library.dump()} @router.post("/get_skills_library_categories") -def get_skills_library_categories(discussionInfos:clientInfos): +def get_skills_library_categories(discussionInfos:ClientInfos): # get_categories returns a list of strings, each entry is a category return {"status":True, "categories":lollmsElfServer.skills_library.get_categories()} @router.post("/get_skills_library_titles_by_category") def get_skills_library_titles(categoryData:CategoryData): + check_access(lollmsElfServer, categoryData.client_id) # Get titles returns a list of dict each entry has id and title return {"status":True, "titles":lollmsElfServer.skills_library.get_titles_by_category(categoryData.category)} @router.post("/get_skills_library_titles") -def get_skills_library_titles(clientInfos:clientInfos): +def get_skills_library_titles(clientInfos:ClientInfos): + check_access(lollmsElfServer, clientInfos.client_id) # Get titles returns a list of dict each entry has id and title return {"status":True, "titles":lollmsElfServer.skills_library.get_titles()} @router.post("/get_skills_library_content") def get_skills_library_content(skillInfos:SkillInfos): + check_access(lollmsElfServer, skillInfos.client_id) # Get the content of the skill from the id, the output is a list of dicts each entry has id, category, title and content return {"status":True, "contents":lollmsElfServer.skills_library.get_skill(skillInfos.skill_id)} +@router.post("/delete_skill") +def delete_skill(delSkillInfos:DeleteSkillInfos): + check_access(lollmsElfServer, delSkillInfos.client_id) + lollmsElfServer.skills_library.remove_entry(delSkillInfos.skill_id) + return {"status":True} + @router.post("/edit_skill") def edit_skill(skillInfos:SkillUpdateInfos): lollmsElfServer.skills_library.update_skill(skillInfos.skill_id, skillInfos.category, skillInfos.title, skillInfos.content) return {"status":True} @router.post("/add_discussion_to_skills_library") -def add_discussion_to_skills_library(discussionInfos:clientInfos): +def add_discussion_to_skills_library(discussionInfos:ClientInfos): lollmsElfServer.ShowBlockingMessage("Learning...") try: client = lollmsElfServer.session.get_client(discussionInfos.client_id) diff --git a/lollms/server/endpoints/lollms_user.py b/lollms/server/endpoints/lollms_user.py index dbcddf1..4f21325 100644 --- a/lollms/server/endpoints/lollms_user.py +++ b/lollms/server/endpoints/lollms_user.py @@ -16,6 +16,7 @@ from lollms.main_config import BaseConfig from lollms.utilities import detect_antiprompt, remove_text_from_string from ascii_colors import ASCIIColors from lollms.databases.discussions_database import DiscussionsDB +from lollms.security import check_access from pathlib import Path from safe_store.text_vectorizer import TextVectorizer, VectorizationMethod, VisualizationMethod import tqdm @@ -26,6 +27,7 @@ import os from PIL import Image class PersonalPathParameters(BaseModel): + client_id:str path:str # ----------------------- Defining router and main class ------------------------------ @@ -33,8 +35,9 @@ class PersonalPathParameters(BaseModel): router = APIRouter() lollmsElfServer = LOLLMSWebUI.get_instance() -@router.get("/switch_personal_path") +@router.post("/switch_personal_path") def switch_personal_path(data:PersonalPathParameters): + client = check_access(lollmsElfServer, data.client_id) path = data.path global_paths_cfg = Path("./global_paths_cfg.yaml") if global_paths_cfg.exists():