Vulnerable code sample

from fastapi import APIRouter, FastAPI
from pydantic import BaseModel
from pathlib import Path
import importlib.machinery
import importlib.util
import sys

# This code is a representation of the vulnerable logic described in CVE-2024-2356.
# It simulates the necessary components of the lollms-webui application
# to demonstrate the vulnerability.

# --- Start of Simulated Application Environment ---

# Simulate the Pydantic model for the request body
class ReinstallExtensionData(BaseModel):
    name: str

# Simulate the paths configuration object
class LollmsPaths:
    def __init__(self):
        # The legitimate directory for extensions
        self.extensions_zoo_path = Path("extensions")

# Simulate the main server object
class LollmsElfServer:
    def __init__(self):
        self.lollms_paths = LollmsPaths()

# Simulate the ExtensionBuilder class that loads the extension code
class ExtensionBuilder:
    def build_extension(self, extension_path, lollms_elf_server):
        # This function attempts to load an __init__.py file from the given path
        init_file = Path(extension_path) / "__init__.py"
        if not init_file.exists():
            # In a real scenario, this might raise an error or log a message.
            return None

        try:
            # The use of SourceFileLoader on a user-controllable path is what
            # enables the execution of arbitrary Python code.
            module_name = str(init_file.stem)
            loader = importlib.machinery.SourceFileLoader(module_name, str(init_file))
            spec = importlib.util.spec_from_loader(loader.name, loader)
            module = importlib.util.module_from_spec(spec)
            
            # This line executes the code within the __init__.py file,
            # leading to Remote Code Execution (RCE).
            loader.exec_module(module)
            
            return module
        except Exception as e:
            # In a real scenario, this might log the error.
            print(f"Error building extension: {e}", file=sys.stderr)
            return None

# --- End of Simulated Application Environment ---


# Instantiate the main application and router objects
app = FastAPI()
router = APIRouter()

# Instantiate a mock server object for the demonstration
lollmsElfServer = LollmsElfServer()


@router.post("/reinstall_extension")
async def reinstall_extension(data: ReinstallExtensionData):
    """
    This is the vulnerable endpoint. It allows reinstalling an extension.
    The 'name' parameter is not sanitized, leading to a Local File Inclusion.
    """
    try:
        # VULNERABLE CODE:
        # The user-controlled 'data.name' is directly concatenated with a base path.
        # An attacker can supply path traversal characters (e.g., '../../')
        # to navigate to other directories on the filesystem.
        extension_path = lollmsElfServer.lollms_paths.extensions_zoo_path / data.name

        # The application then attempts to build/load the extension from the
        # constructed path. If the path points to a directory containing a
        # malicious '__init__.py' file (e.g., in an upload directory),
        # that file will be executed.
        ExtensionBuilder().build_extension(extension_path, lollmsElfServer)

        return {"status": "success", "message": f"Extension {data.name} reinstalled."}
    except Exception as ex:
        return {"status": "failure", "error": str(ex)}


app.include_router(router)