Patched code sample

import dataclasses
from typing import Any, Dict, Optional


# Mock HTTPException for demonstration purposes
class HTTPException(Exception):
    def __init__(self, status_code: int, detail: str):
        self.status_code = status_code
        self.detail = detail
        super().__init__(f"{status_code}: {detail}")


# A simplified representation of a Block, based on the vulnerability description
@dataclasses.dataclass
class Block:
    uuid: str
    name: str
    disabled: bool

    def execute(self, *args, **kwargs) -> Any:
        """Simulates the execution of a block."""
        print(f"Executing block '{self.name}' (UUID: {self.uuid})...")
        # In a real scenario, this would perform the block's specific logic.
        # For the vulnerable 'BlockInstallationBlock', this is where arbitrary
        # code would be written to a file and imported.
        if self.name == "BlockInstallationBlock":
            print("SIMULATING DANGEROUS ACTION: Writing and importing arbitrary code.")
            # pass # In a real exploit, this would be os.system() or __import__()
        return {"status": "success", "output": f"Block '{self.name}' executed."}


# Mock database of installed blocks
# The key vulnerability context is that 'BlockInstallationBlock' is disabled
# but could still be executed in the vulnerable version.
MOCK_BLOCK_DB: Dict[str, Block] = {
    "uuid-safe-block-123": Block(
        uuid="uuid-safe-block-123", name="DataProcessingBlock", disabled=False
    ),
    "uuid-vulnerable-block-456": Block(
        uuid="uuid-vulnerable-block-456",
        name="BlockInstallationBlock",
        disabled=True,  # This block should NOT be executable
    ),
}


def find_block_by_uuid(uuid: str) -> Optional[Block]:
    """Simulates fetching a block's configuration from a database."""
    return MOCK_BLOCK_DB.get(uuid)


# This function represents the fixed API endpoint for block execution.
def execute_block(block_uuid: str) -> Dict[str, Any]:
    """
    Finds and executes a block by its UUID, but only if it is not disabled.
    This demonstrates the fix for CVE-2024-27880.
    """
    block = find_block_by_uuid(block_uuid)

    if not block:
        raise HTTPException(status_code=404, detail=f"Block with UUID '{block_uuid}' not found.")

    # =========================================================================
    # THE FIX:
    # The vulnerability was a failure to check the `disabled` flag. The fix
    # is to add this check and prevent execution if the block is disabled.
    if block.disabled:
        raise HTTPException(
            status_code=403,
            detail=f"Execution of disabled block '{block.name}' is not allowed."
        )
    # =========================================================================

    # If the block is not disabled, proceed with execution
    result = block.execute()
    return result


if __name__ == "__main__":
    print("--- DEMONSTRATING THE FIX ---")

    # 1. Attempt to execute a normal, enabled block. This should succeed.
    print("\nAttempting to execute an enabled block...")
    try:
        enabled_block_uuid = "uuid-safe-block-123"
        response = execute_block(enabled_block_uuid)
        print(f"SUCCESS: API response: {response}")
    except HTTPException as e:
        print(f"ERROR: Caught exception: {e}")

    # 2. Attempt to execute the disabled 'BlockInstallationBlock'.
    #    This should be blocked by the fix.
    print("\nAttempting to execute a disabled block (the exploit target)...")
    try:
        disabled_block_uuid = "uuid-vulnerable-block-456"
        response = execute_block(disabled_block_uuid)
        print(f"ERROR: API response: {response}")
    except HTTPException as e:
        print(f"SUCCESS (FIX WORKED): Caught expected exception: {e}")