import demistomock as demisto # noqa: F401 from CommonServerPython import * # noqa: F401 import json import os import shutil """ CLIENT CLASS """ class Client(BaseClient): def test_module(self): self._http_request("GET", "/v1/auth") def set_airs_client(self, airs_client): self.airs_client = airs_client def document_list(self): return self._list("documents") def document_get(self, folder: str, document: str): try: name = document_name(folder, document, self.document_list()) response = self._http_request(method="GET", url_suffix=f"/v1/document/{name}") except Exception as e: msg = f"AnythingLLM: document_get: exception getting document details - {e}" demisto.debug(msg) raise Exception(msg) return response def document_delete(self, folder: str, document: str): try: name = document_name(folder, document, self.document_list()) data = {"names": [f"{folder}/{name}"]} response = self._http_request(method="DELETE", url_suffix="/v1/system/remove-documents", json_data=data) except Exception as e: msg = f"AnythingLLM: document_delete: exception deleting document - {e}" demisto.debug(msg) raise Exception(msg) return {"message": response} def document_createfolder(self, folder: str): try: data = {"name": folder} response = self._http_request(method="POST", url_suffix="/v1/document/create-folder", json_data=data) except Exception as e: msg = f"AnythingLLM: document_createfolder: exception creating folder - {e}" demisto.debug(msg) raise Exception(msg) return response def document_move(self, srcfolder: str, dstfolder: str, document: str): try: name = document_name(srcfolder, document, self.document_list()) data = {"files": [{"from": f"{srcfolder}/{name}", "to": f"{dstfolder}/{name}"}]} response = self._http_request(method="POST", url_suffix="/v1/document/move-files", json_data=data) except Exception as e: msg = f"AnythingLLM: document_move: exception moving document - {e}" demisto.debug(msg) raise Exception(msg) return response def document_upload_text(self, text: str, title: str, description: str, author: str, source: str): try: try: exists = False document_name("custom-documents", title, self.document_list()) exists = True except Exception: data = { "textContent": text, "metadata": {"title": title, "docAuthor": author, "description": description, "docSource": source}, } response = self._http_request(method="POST", url_suffix="/v1/document/raw-text", json_data=data) finally: if exists: # pylint: disable=E0601 raise Exception(f"document already exists [{title}]") except Exception as e: msg = f"AnythingLLM: document_upload_text: exception uploading text - {e}" demisto.debug(msg) raise Exception(msg) return response # pylint: disable=E0601 def document_upload_link(self, link: str, title: str, description: str, author: str, source: str): try: try: exists = False document_name("custom-documents", title, self.document_list()) exists = True except Exception: data = { "link": link, "metadata": {"title": title, "docAuthor": author, "description": description, "docSource": source}, } response = self._http_request(method="POST", url_suffix="/v1/document/upload-link", json_data=data) finally: if exists: # pylint: disable=E0601 raise Exception(f"document already exists [{title}]") except Exception as e: msg = f"AnythingLLM: document_upload_link: exception uploading link [{link}] - {e}" demisto.debug(msg) raise Exception(msg) return response # pylint: disable=E0601 def document_upload_file(self, entry_id): file_name = "" try: headers = self._headers del headers["Content-Type"] file_path = demisto.getFilePath(entry_id)["path"] file_name = os.path.basename(demisto.getFilePath(entry_id)["name"]) try: exists = False document_name("custom-documents", file_name, self.document_list()) exists = True except Exception: shutil.copy(file_path, file_name) response = self._http_request( method="POST", headers=headers, url_suffix="/v1/document/upload", files={"file": (f"{entry_id}_{file_name}", open(file_name, "rb"))}, ) finally: if exists: # pylint: disable=E0601 raise Exception(f"document already exists [{file_name}]") except Exception as e: msg = f"AnythingLLM: document_upload_file: exception uploading a file entry [{entry_id}] from the war room - {e}" demisto.debug(msg) raise Exception(msg) finally: try: if file_name: os.remove(file_name) except OSError: pass return response # pylint: disable=E0601 def workspace_new(self, workspace: str): try: if len(workspace.strip()) == 0: raise Exception("workspace parameter is blank") try: exists = False workspace_slug(workspace, self.workspace_list()) exists = True except Exception: data = {"name": workspace} response = self._http_request(method="POST", url_suffix="/v1/workspace/new", json_data=data) return response finally: if exists: # pylint: disable=E0601 raise Exception("workspace already exists") except Exception as e: msg = f"AnythingLLM: workspace_new: exception creating a new workspace [{workspace}] - {e}" demisto.debug(msg) raise Exception(msg) def workspace_chat(self, workspace: str, message: str, mode: str): return self._chat(workspace, message, mode, "chat") def workspace_stream_chat(self, workspace: str, message: str, mode: str): return self._chat(workspace, message, mode, "stream-chat") def workspace_list(self): return self._list("workspaces") def thread_list(self, workspace: str): try: slug = workspace_slug(workspace, self.workspace_list()) response = self._http_request( method="GET", url_suffix=f"/v1/workspace/{slug}", ) except Exception as e: msg = f"AnythingLLM: thread_list: exception listing workspace threads - {e}" demisto.debug(msg) raise Exception(msg) return response["workspace"][0]["threads"] def workspace_get(self, workspace: str): try: slug = workspace_slug(workspace, self.workspace_list()) response = self._http_request( method="GET", url_suffix=f"/v1/workspace/{slug}", ) except Exception as e: msg = f"AnythingLLM: workspace_get: exception getting workspace details - {e}" demisto.debug(msg) raise Exception(msg) return response def workspace_delete(self, workspace: str): try: slug = workspace_slug(workspace, self.workspace_list()) self._http_request(method="DELETE", url_suffix=f"/v1/workspace/{slug}", resp_type="bytes") except Exception as e: msg = f"AnythingLLM: workspace_delete: exception deleting workspace - {e}" demisto.debug(msg) raise Exception(msg) return {"message": {"success": True, "message": "Workspace removed successfully"}} def workspace_thread_new(self, workspace: str, thread: str): try: wslug = workspace_slug(workspace, self.workspace_list()) tslug = f"{thread}_slug" threads = self.thread_list(workspace) if any(t["slug"] == tslug for t in threads): raise Exception(f"Thread '{thread}' already exists in workspace '{workspace}'.") data = {"name": thread, "slug": tslug} response = self._http_request(method="POST", url_suffix=f"/v1/workspace/{wslug}/thread/new", json_data=data) return response except Exception as e: msg = f"AnythingLLM: workspace_thread_new: exception creating a new workspace [{workspace}] - {e}" demisto.debug(msg) raise Exception(msg) def workspace_thread_chat(self, workspace: str, thread: str, message: str, mode: str): if demisto.params().get("airs_scan_prompt"): result = self.airs_client.airs_sync_scan_prompt(message) if result["action"] == "block": msg = f"AnythingLLM: workspace_thread_chat: prompt blocked [{result.get('prompt_detected')}]" demisto.debug(msg) raise Exception(msg) response = self._tchat(workspace, thread, message, mode, "chat") if demisto.params().get("airs_scan_response"): result = self.airs_client.airs_sync_scan_response(response["textResponse"]) if result["action"] == "block": msg = f"AnythingLLM: workspace_thread_chat: response blocked [{result.get('response_detected')}]" demisto.debug(msg) raise Exception(msg) return response def workspace_thread_chats(self, workspace: str, thread: str): try: slug = workspace_slug(workspace, self.workspace_list()) response = self._http_request(method="GET", url_suffix=f"/v1/workspace/{slug}/thread/{thread + '_slug'}/chats") except Exception as e: msg = f"AnythingLLM: workspace_thread_chats: exception chatting - {e}" demisto.debug(msg) raise Exception(msg) return response def workspace_settings(self, workspace: str, settings: dict): try: settings = validate_workspace_settings(settings) if len(settings) == 0: raise Exception("Invalid workspace settings") slug = workspace_slug(workspace, self.workspace_list()) response = self._http_request(method="POST", url_suffix=f"/v1/workspace/{slug}/update", json_data=settings) except Exception as e: msg = f"AnythingLLM: workspace_settings: exception updating workspace settings - {e}" demisto.debug(msg) raise Exception(msg) return response def workspace_thread_delete(self, workspace: str, thread: str): try: wslug = workspace_slug(workspace, self.workspace_list()) tslug = thread_slug(thread, self.thread_list(workspace)) self._http_request(method="DELETE", url_suffix=f"/v1/workspace/{wslug}/thread/{tslug}", resp_type="bytes") except Exception as e: msg = f"AnythingLLM: workspace_thread_delete: exception deleting workspace - {e}" demisto.debug(msg) raise Exception(msg) return {"message": {"success": True, "message": "Conversation thread removed successfully"}} def workspace_add_embedding(self, workspace: str, folder: str, document: str): return self._embedding(workspace, folder, document, "adds") def workspace_delete_embedding(self, workspace: str, folder: str, document: str): return self._embedding(workspace, folder, document, "deletes") def workspace_pin(self, workspace: str, folder: str, document: str, status: str): try: if status.lower() == "true": pinst = True elif status.lower() == "false": pinst = False else: raise Exception("document pin status of [true] or [false] not passed") name = document_name(folder, document, self.document_list()) data = {"docPath": f"{folder}/{name}", "pinStatus": pinst} slug = workspace_slug(workspace, self.workspace_list()) response = self._http_request(method="POST", url_suffix=f"/v1/workspace/{slug}/update-pin", json_data=data) except Exception as e: msg = f"AnythingLLM: workspace_pin: exception pinning embedded document to workspace - {e}" demisto.debug(msg) raise Exception(msg) return response def _chat(self, workspace: str, message: str, mode: str, ttype: str): try: data = {"message": message, "mode": validate_chat_mode(mode)} slug = workspace_slug(workspace, self.workspace_list()) if demisto.params().get("airs_scan_prompt"): result = self.airs_client.airs_sync_scan_prompt(message) if result["action"] == "block": msg = f"AnythingLLM: _chat: prompt blocked [{result.get('prompt_detected')}]" demisto.debug(msg) raise Exception(msg) response = self._http_request(method="POST", url_suffix=f"/v1/workspace/{slug}/{ttype}", json_data=data) if demisto.params().get("airs_scan_response"): result = self.airs_client.airs_sync_scan_response(response.get("textResponse")) if result["action"] == "block": msg = f"AnythingLLM: _chat: response blocked [{result.get('response_detected')}]" demisto.debug(msg) raise Exception(msg) except Exception as e: msg = f"AnythingLLM: _chat: exception chatting - {e}" demisto.debug(msg) raise Exception(msg) return response def _tchat(self, workspace: str, thread: str, message: str, mode: str, ttype: str): try: data = {"message": message, "mode": validate_chat_mode(mode)} wslug = workspace_slug(workspace, self.workspace_list()) tslug = thread_slug(thread, self.thread_list(workspace)) response = self._http_request( method="POST", url_suffix=f"/v1/workspace/{wslug}/thread/{tslug}/{ttype}", json_data=data ) except Exception as e: msg = f"AnythingLLM: _tchat: exception chatting - {e}" demisto.debug(msg) raise Exception(msg) return response def _list(self, items: str): try: response = self._http_request( method="GET", url_suffix=f"/v1/{items}", ) except Exception as e: msg = f"AnythingLLM: _list: exception listing {items} - {e}" demisto.debug(msg) raise Exception(msg) return response def _embedding(self, workspace: str, folder: str, document: str, action: str): try: name = "" name = document_name(folder, document, self.document_list()) try: ws = self.workspace_get(workspace) except Exception: raise Exception("workspace not found") if action == "adds": if embedding_exists(ws, document): raise Exception("already embedded") elif action == "deletes": if not embedding_exists(ws, document): raise Exception("not embedded") else: raise Exception(f"action [{action}] not 'adds' or 'deletes' ") data = {action: [f"{folder}/{name}"]} slug = workspace_slug(workspace, self.workspace_list()) response = self._http_request(method="POST", url_suffix=f"/v1/workspace/{slug}/update-embeddings", json_data=data) except Exception as e: msg = f"AnythingLLM: _embedding: exception [{action}] a document embedding [{document}] in [{workspace}] - {e}" demisto.debug(msg) raise Exception(msg) return response class AirsClient(BaseClient): def test_module(self): self.airs_sync_scan_prompt("Hello!") def set_params(self, params: dict): self.params = params def airs_sync_scan_prompt(self, message: str) -> dict: new_args = { "profile_name": self.params.get("airs_profile"), "code_prompt": self.params.get("code_prompt"), "code_response": self.params.get("code_response"), "prompt": message, "response": self.params.get("response"), "context": self.params.get("context"), "ai_model": self.params.get("airs_model"), "app_name": self.params.get("airs_app"), "app_user": self.params.get("airs_user"), "user_ip": "", } return self.airs_sync_scan(new_args) def airs_sync_scan_response(self, message: str) -> dict: new_args = { "profile_name": self.params.get("airs_profile"), "code_prompt": self.params.get("code_prompt"), "code_response": self.params.get("code_response"), "prompt": self.params.get("prompt"), "response": message, "context": self.params.get("context"), "ai_model": self.params.get("airs_model"), "app_name": self.params.get("airs_app"), "app_user": self.params.get("airs_user"), "user_ip": "", } return self.airs_sync_scan(new_args) def airs_sync_scan(self, args: dict) -> dict: data = { "ai_profile": {"profile_name": args.get("profile_name")}, "contents": [ { "code_prompt": args.get("code_prompt"), "code_response": args.get("code_response"), "prompt": args.get("prompt"), "response": args.get("response"), "context": args.get("context"), } ], "metadata": { "ai_model": args.get("ai_model"), "app_name": args.get("app_name"), "app_user": args.get("app_user"), "user_ip": args.get("user_ip"), }, "tr_id": "", } headers = self._headers headers["Content-Type"] = "application/json" response = self._http_request("POST", "v1/scan/sync/request", json_data=data, headers=headers) return response """ HELPER FUNCTIONS """ def embedding_exists(ws: dict, document: str) -> bool: if "documents" in ws["workspace"][0]: for doc in ws["workspace"][0]["documents"]: metadata = json.loads(doc["metadata"]) if metadata["title"] == document: return True return False def workspace_slug(workspace: str, workspaces) -> str: for w in workspaces["workspaces"]: if w["name"] == workspace: return w["slug"] raise Exception(f"workspace name not found [{workspace}]") def thread_slug(thread: str, threads: list) -> str: for t in threads: # Thread data returned does not include the name, so for now, # enforce unique thread names and always append "_slug" to the name # to create the slug if t["slug"] == thread + "_slug": return t["slug"] raise Exception(f"thread name not found [{thread}]") def normal_document_title(title: str) -> str: title = " ".join(title.strip().split()) return title.lower().replace(" ", "-") # + ".txt" def remove_entryid(title: str) -> str: parts = title.split("_", 1) # Strip the entry_id when file is uploaded from XSOAR 225@12345_