wehub-resource-sync
84 lines
2.7 KiB
Python
84 lines
2.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Standard
|
|
from typing import Any, List
|
|
import importlib
|
|
|
|
# Third Party
|
|
from fastapi import APIRouter
|
|
from starlette.requests import Request
|
|
from starlette.responses import PlainTextResponse
|
|
|
|
# First Party
|
|
from lmcache.logging import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
def _get_allowed_imports(request: Request) -> List[str]:
|
|
"""Extract script_allowed_imports from either inProcess or mp mode."""
|
|
adapter = getattr(request.app.state, "lmcache_adapter", None)
|
|
if adapter is not None:
|
|
return getattr(adapter.config, "script_allowed_imports", None) or []
|
|
configs = getattr(request.app.state, "configs", None)
|
|
if isinstance(configs, dict):
|
|
mp_cfg = configs.get("mp")
|
|
return getattr(mp_cfg, "script_allowed_imports", None) or []
|
|
return []
|
|
|
|
|
|
@router.post("/run_script")
|
|
async def run_script(request: Request):
|
|
form_data = await request.form()
|
|
script_file = form_data.get("script")
|
|
|
|
if not script_file or not hasattr(script_file, "file"):
|
|
return PlainTextResponse("No script file provided", status_code=400)
|
|
|
|
script_content = await script_file.read()
|
|
|
|
try:
|
|
allowed_imports = _get_allowed_imports(request)
|
|
|
|
# Pre-import allowed modules
|
|
allowed_modules = {}
|
|
for module_name in allowed_imports:
|
|
try:
|
|
module = importlib.import_module(module_name)
|
|
allowed_modules[module_name] = module
|
|
logger.info("Imported allowed module: %s", module_name)
|
|
except ImportError as e:
|
|
logger.warning("Failed to import module %s: %s", module_name, e)
|
|
|
|
# Create custom __import__ function that only allows configured modules
|
|
def restricted_import(name, globals=None, locals=None, fromlist=(), level=0):
|
|
if name in allowed_modules:
|
|
return allowed_modules[name]
|
|
raise ImportError(f"Import of '{name}' is not allowed")
|
|
|
|
restricted_globals = {
|
|
"__builtins__": {
|
|
"print": print,
|
|
"str": str,
|
|
"int": int,
|
|
"float": float,
|
|
"list": list,
|
|
"dict": dict,
|
|
"tuple": tuple,
|
|
"set": set,
|
|
"__import__": restricted_import,
|
|
},
|
|
"app": request.app,
|
|
}
|
|
|
|
restricted_locals: dict[str, Any] = {}
|
|
|
|
exec(script_content, restricted_globals, restricted_locals)
|
|
|
|
result = restricted_locals.get("result", "Script executed successfully")
|
|
return PlainTextResponse(str(result), media_type="text/plain")
|
|
|
|
except Exception as e:
|
|
return PlainTextResponse(f"Error executing script: {str(e)}", status_code=500)
|