mirror of
https://github.com/open-webui/open-webui.git
synced 2025-12-12 04:15:25 +00:00
Merge pull request #14703 from rragundez/code-interpreter-blacklist
feat: Blacklist modules from arbitrary code execution in code interpreter
This commit is contained in:
commit
47560d4d72
2 changed files with 31 additions and 0 deletions
|
|
@ -1857,6 +1857,16 @@ CODE_INTERPRETER_JUPYTER_TIMEOUT = PersistentConfig(
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
CODE_INTERPRETER_BLACKLISTED_MODULES = PersistentConfig(
|
||||||
|
"CODE_INTERPRETER_BLACKLISTED_MODULES",
|
||||||
|
"code_interpreter.blacklisted_modules",
|
||||||
|
[
|
||||||
|
library.strip()
|
||||||
|
for library in os.environ.get("CODE_INTERPRETER_BLACKLISTED_MODULES", "").split(",")
|
||||||
|
if library.strip()
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
DEFAULT_CODE_INTERPRETER_PROMPT = """
|
DEFAULT_CODE_INTERPRETER_PROMPT = """
|
||||||
#### Tools Available
|
#### Tools Available
|
||||||
|
|
|
||||||
|
|
@ -3,6 +3,7 @@ import logging
|
||||||
import sys
|
import sys
|
||||||
import os
|
import os
|
||||||
import base64
|
import base64
|
||||||
|
import textwrap
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
from aiocache import cached
|
from aiocache import cached
|
||||||
|
|
@ -91,6 +92,7 @@ from open_webui.config import (
|
||||||
CACHE_DIR,
|
CACHE_DIR,
|
||||||
DEFAULT_TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE,
|
DEFAULT_TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE,
|
||||||
DEFAULT_CODE_INTERPRETER_PROMPT,
|
DEFAULT_CODE_INTERPRETER_PROMPT,
|
||||||
|
CODE_INTERPRETER_BLACKLISTED_MODULES,
|
||||||
)
|
)
|
||||||
from open_webui.env import (
|
from open_webui.env import (
|
||||||
SRC_LOG_LEVELS,
|
SRC_LOG_LEVELS,
|
||||||
|
|
@ -2369,6 +2371,25 @@ async def process_chat_response(
|
||||||
try:
|
try:
|
||||||
if content_blocks[-1]["attributes"].get("type") == "code":
|
if content_blocks[-1]["attributes"].get("type") == "code":
|
||||||
code = content_blocks[-1]["content"]
|
code = content_blocks[-1]["content"]
|
||||||
|
if CODE_INTERPRETER_BLACKLISTED_MODULES:
|
||||||
|
blocking_code = textwrap.dedent(f"""
|
||||||
|
import builtins
|
||||||
|
|
||||||
|
BLACKLISTED_MODULES = {CODE_INTERPRETER_BLACKLISTED_MODULES}
|
||||||
|
|
||||||
|
_real_import = builtins.__import__
|
||||||
|
def restricted_import(name, globals=None, locals=None, fromlist=(), level=0):
|
||||||
|
if name.split('.')[0] in BLACKLISTED_MODULES:
|
||||||
|
importer_name = globals.get('__name__') if globals else None
|
||||||
|
if importer_name == '__main__':
|
||||||
|
raise ImportError(
|
||||||
|
f"Direct import of module {{name}} is restricted."
|
||||||
|
)
|
||||||
|
return _real_import(name, globals, locals, fromlist, level)
|
||||||
|
|
||||||
|
builtins.__import__ = restricted_import
|
||||||
|
""")
|
||||||
|
code = blocking_code + "\n" + code
|
||||||
|
|
||||||
if (
|
if (
|
||||||
request.app.state.config.CODE_INTERPRETER_ENGINE
|
request.app.state.config.CODE_INTERPRETER_ENGINE
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue