mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-08-28 05:27:24 +00:00
feat: add YOLO Mode for Jupyter notebooks
- Add comprehensive security validation for Jupyter notebooks - Implement YOLO Mode setting to bypass security restrictions in trusted environments - Add jupyterNotebookYoloMode configuration option - Create JupyterNotebookSecurity class with risk detection for: - Critical risks: eval, exec, shell commands - High risks: dangerous imports, script injection - Medium risks: file/network operations - Implement automatic sanitization of dangerous content - Add read-only mode enforcement for risky notebooks - Include comprehensive test suite with 27 test cases - Add detailed documentation for security features This allows users to work with trusted notebooks containing "unsafe" constructs by enabling YOLO Mode in settings when needed.
This commit is contained in:
Roo Code
parent
b486282f6b
commit
24b0f76ff1
8 changed files with 1468 additions and 13 deletions
296
docs/jupyter-notebook-security.md
Normal file
296
docs/jupyter-notebook-security.md
Normal file
|
|
@ -0,0 +1,296 @@
|
|||
# Jupyter Notebook Security Features
|
||||
|
||||
## Overview
|
||||
|
||||
Roo Code includes comprehensive security features for working with Jupyter notebooks (`.ipynb` files). These features protect against potentially dangerous code execution, cross-site scripting (XSS) attacks, and other security risks commonly associated with Jupyter notebooks.
|
||||
|
||||
## Security Validation
|
||||
|
||||
### Automatic Risk Detection
|
||||
|
||||
When opening or editing Jupyter notebooks, Roo Code automatically scans for:
|
||||
|
||||
#### Critical Risks (Blocks Editing)
|
||||
|
||||
- **Code Execution**: `eval()`, `exec()`, `compile()`, `__import__()`
|
||||
- **Shell Commands**: `!command`, `%system`, `%%bash`, `%%sh`, `%%script`
|
||||
|
||||
#### High Risks (Blocks Editing)
|
||||
|
||||
- **Dangerous Imports**: `subprocess`, `os`, `sys`, `socket`, `pickle`
|
||||
- **Script Injection**: `<script>` tags in markdown cells
|
||||
- **IFrame Injection**: `<iframe>` tags in markdown cells
|
||||
- **JavaScript Protocol**: `javascript:` URLs in markdown
|
||||
- **JavaScript Outputs**: Cells with `application/javascript` output
|
||||
|
||||
#### Medium Risks (Warnings Only)
|
||||
|
||||
- **File Operations**: `open()`, `file()`, `.read()`, `.write()`
|
||||
- **Network Operations**: `requests`, `urllib`, `http.client`
|
||||
- **Event Handlers**: `onclick`, `onload`, etc. in markdown
|
||||
- **Oversized Cells**: Cells exceeding 1MB
|
||||
- **Excessive Cell Count**: Notebooks with >1000 cells
|
||||
|
||||
#### Low Risks (Informational)
|
||||
|
||||
- **Suspicious Metadata**: Presence of `widgets`, `extensions`, or `jupyter_dashboards` fields
|
||||
|
||||
### Read-Only Mode
|
||||
|
||||
Notebooks containing critical or high-severity risks automatically open in read-only mode, preventing:
|
||||
|
||||
- Cell modifications
|
||||
- Cell additions or deletions
|
||||
- Saving changes to disk
|
||||
|
||||
## YOLO Mode (You Only Live Once)
|
||||
|
||||
For trusted environments where you need to work with notebooks containing "unsafe" constructs, YOLO Mode allows you to bypass all security restrictions.
|
||||
|
||||
### Enabling YOLO Mode
|
||||
|
||||
Add the following to your VSCode settings:
|
||||
|
||||
```json
|
||||
{
|
||||
"roo-code.jupyterNotebookYoloMode": true
|
||||
}
|
||||
```
|
||||
|
||||
Or use the command palette:
|
||||
|
||||
1. Press `Cmd/Ctrl + Shift + P`
|
||||
2. Search for "Preferences: Open Settings (JSON)"
|
||||
3. Add the setting above
|
||||
|
||||
### What YOLO Mode Does
|
||||
|
||||
When enabled, YOLO Mode:
|
||||
|
||||
- ✅ Bypasses all security validation
|
||||
- ✅ Allows editing of any notebook content
|
||||
- ✅ Permits execution of shell commands
|
||||
- ✅ Allows dangerous imports and operations
|
||||
- ✅ Enables saving of modified notebooks
|
||||
|
||||
### ⚠️ YOLO Mode Warnings
|
||||
|
||||
**USE WITH EXTREME CAUTION**: YOLO Mode completely disables security protections. Only enable it when:
|
||||
|
||||
- Working with trusted notebooks from known sources
|
||||
- In isolated/sandboxed environments
|
||||
- You fully understand the risks involved
|
||||
- The notebook requires system-level operations for legitimate purposes
|
||||
|
||||
## Security Configuration
|
||||
|
||||
### Advanced Configuration Options
|
||||
|
||||
```typescript
|
||||
interface JupyterSecurityConfig {
|
||||
// Enable YOLO Mode (bypass all security)
|
||||
yoloMode?: boolean
|
||||
|
||||
// Allow code execution (default: false)
|
||||
allowCodeExecution?: boolean
|
||||
|
||||
// Force read-only mode (default: true)
|
||||
readOnlyMode?: boolean
|
||||
|
||||
// Maximum cell size in bytes (default: 1MB)
|
||||
maxCellSize?: number
|
||||
|
||||
// Maximum number of cells (default: 1000)
|
||||
maxCellCount?: number
|
||||
|
||||
// Trusted source paths (notebooks from these paths bypass security)
|
||||
trustedSources?: string[]
|
||||
}
|
||||
```
|
||||
|
||||
### Trusted Sources
|
||||
|
||||
You can configure specific directories as trusted sources. Notebooks from these locations will bypass security checks:
|
||||
|
||||
```json
|
||||
{
|
||||
"roo-code.jupyterTrustedSources": ["/home/user/trusted-notebooks", "/projects/data-science/verified"]
|
||||
}
|
||||
```
|
||||
|
||||
## Automatic Sanitization
|
||||
|
||||
When security risks are detected and YOLO Mode is disabled, Roo Code automatically sanitizes dangerous content:
|
||||
|
||||
### Code Cell Sanitization
|
||||
|
||||
- Dangerous code is commented out with security warnings
|
||||
- Original code is preserved as comments for reference
|
||||
- Cell outputs are cleared
|
||||
- Execution count is reset
|
||||
|
||||
Example:
|
||||
|
||||
```python
|
||||
# ⚠️ SECURITY WARNING: This cell has been disabled due to security risks (eval, exec)
|
||||
# To run this cell, enable YOLO Mode in settings
|
||||
|
||||
# Original code:
|
||||
# eval('__import__("os").system("rm -rf /")')
|
||||
```
|
||||
|
||||
### Markdown Cell Sanitization
|
||||
|
||||
- Script tags are replaced with `[REMOVED: script tag]`
|
||||
- IFrame tags are replaced with `[REMOVED: iframe]`
|
||||
- JavaScript protocols are replaced with `[REMOVED]:`
|
||||
- Event handlers are prefixed with `data-removed=`
|
||||
|
||||
### Output Sanitization
|
||||
|
||||
- HTML outputs containing scripts are removed
|
||||
- JavaScript outputs are replaced with plain text warnings
|
||||
- Data URIs with HTML content are flagged
|
||||
|
||||
## Best Practices
|
||||
|
||||
### For Notebook Authors
|
||||
|
||||
1. **Avoid dangerous patterns**: Use subprocess.run() instead of eval/exec
|
||||
2. **Validate inputs**: Never execute user input directly
|
||||
3. **Use virtual environments**: Isolate notebook dependencies
|
||||
4. **Document requirements**: Clearly state what system access is needed
|
||||
5. **Test in sandboxes**: Verify notebooks work in restricted environments
|
||||
|
||||
### For Notebook Users
|
||||
|
||||
1. **Review before running**: Always inspect notebook content before execution
|
||||
2. **Use YOLO Mode sparingly**: Only for trusted, verified notebooks
|
||||
3. **Work in containers**: Use Docker or VMs for untrusted notebooks
|
||||
4. **Check outputs**: Review cell outputs for suspicious content
|
||||
5. **Keep backups**: Save important data before running unknown notebooks
|
||||
|
||||
## Security Risk Examples
|
||||
|
||||
### Critical Risk Example
|
||||
|
||||
```python
|
||||
# This will be blocked
|
||||
eval('__import__("os").system("cat /etc/passwd")')
|
||||
exec(user_input) # Never execute user input
|
||||
!rm -rf / # Shell commands are blocked
|
||||
```
|
||||
|
||||
### High Risk Example
|
||||
|
||||
```python
|
||||
# These imports will trigger read-only mode
|
||||
import subprocess
|
||||
import os
|
||||
import socket
|
||||
|
||||
# Markdown with scripts
|
||||
"""
|
||||
<script>alert('XSS')</script>
|
||||
<iframe src="http://malicious.site"></iframe>
|
||||
"""
|
||||
```
|
||||
|
||||
### Medium Risk Example
|
||||
|
||||
```python
|
||||
# These will show warnings but allow editing
|
||||
with open('/etc/hosts', 'r') as f:
|
||||
data = f.read()
|
||||
|
||||
import requests
|
||||
response = requests.get('https://api.example.com')
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### "Cannot modify notebook: Security risks detected"
|
||||
|
||||
- **Cause**: The notebook contains critical or high-severity security risks
|
||||
- **Solution**: Enable YOLO Mode if you trust the notebook source
|
||||
|
||||
### "Notebook is in read-only mode"
|
||||
|
||||
- **Cause**: Security validation detected dangerous patterns
|
||||
- **Solution**: Review the security warnings and enable YOLO Mode if needed
|
||||
|
||||
### "Cell has been disabled due to security risks"
|
||||
|
||||
- **Cause**: Automatic sanitization has commented out dangerous code
|
||||
- **Solution**: Review the code and enable YOLO Mode to restore functionality
|
||||
|
||||
## API Integration
|
||||
|
||||
### Programmatic Security Validation
|
||||
|
||||
```typescript
|
||||
import { JupyterNotebookSecurity } from "./jupyter-notebook-security"
|
||||
import { JupyterNotebookHandler } from "./jupyter-notebook-handler"
|
||||
|
||||
// Create security validator
|
||||
const security = new JupyterNotebookSecurity({
|
||||
yoloMode: false,
|
||||
maxCellSize: 1024 * 1024,
|
||||
trustedSources: ["/trusted/path"],
|
||||
})
|
||||
|
||||
// Load and validate notebook
|
||||
const handler = await JupyterNotebookHandler.fromFile("notebook.ipynb", { yoloMode: false })
|
||||
|
||||
// Check security status
|
||||
if (handler.hasSecurityRisks()) {
|
||||
const risks = handler.getSecurityRisks()
|
||||
console.log("Security risks detected:", risks)
|
||||
}
|
||||
|
||||
// Check if editing is allowed
|
||||
if (handler.isReadOnly()) {
|
||||
console.log("Notebook is in read-only mode")
|
||||
}
|
||||
```
|
||||
|
||||
### Enabling YOLO Mode Programmatically
|
||||
|
||||
```typescript
|
||||
// Enable YOLO Mode for a specific notebook
|
||||
handler.setYoloMode(true)
|
||||
|
||||
// Now editing is allowed regardless of security risks
|
||||
handler.updateCell(0, 'eval("now this works")')
|
||||
await handler.save()
|
||||
```
|
||||
|
||||
## Security Compliance
|
||||
|
||||
This implementation follows security best practices:
|
||||
|
||||
- **Defense in Depth**: Multiple layers of security validation
|
||||
- **Fail Secure**: Defaults to read-only mode when risks are detected
|
||||
- **Transparency**: Clear warnings about detected risks
|
||||
- **User Control**: YOLO Mode for informed consent
|
||||
- **Preservation**: Original content is preserved during sanitization
|
||||
|
||||
## Future Enhancements
|
||||
|
||||
Planned improvements to the security system:
|
||||
|
||||
- [ ] Configurable risk severity levels
|
||||
- [ ] Custom pattern definitions
|
||||
- [ ] Sandbox execution environment
|
||||
- [ ] Digital signatures for trusted notebooks
|
||||
- [ ] Security audit logs
|
||||
- [ ] Integration with corporate security policies
|
||||
|
||||
## Support
|
||||
|
||||
For questions or issues related to Jupyter notebook security:
|
||||
|
||||
1. Check this documentation
|
||||
2. Review the security warnings in the UI
|
||||
3. Open an issue on GitHub with the `jupyter-security` label
|
||||
4. Contact the security team for enterprise deployments
|
||||
|
|
@ -152,6 +152,9 @@ export const globalSettingsSchema = z.object({
|
|||
hasOpenedModeSelector: z.boolean().optional(),
|
||||
lastModeExportPath: z.string().optional(),
|
||||
lastModeImportPath: z.string().optional(),
|
||||
|
||||
// Jupyter notebook security settings
|
||||
jupyterNotebookYoloMode: z.boolean().optional(),
|
||||
})
|
||||
|
||||
export type GlobalSettings = z.infer<typeof globalSettingsSchema>
|
||||
|
|
|
|||
|
|
@ -1,14 +1,17 @@
|
|||
import { DiffStrategy, DiffResult, ToolUse } from "../../../shared/tools"
|
||||
import { ToolProgressStatus } from "@roo-code/types"
|
||||
import { ToolProgressStatus, GlobalState } from "@roo-code/types"
|
||||
import { JupyterNotebookHandler } from "../../../integrations/misc/jupyter-notebook-handler"
|
||||
import { JupyterSecurityConfig } from "../../../integrations/misc/jupyter-notebook-security"
|
||||
import { MultiSearchReplaceDiffStrategy } from "./multi-search-replace"
|
||||
|
||||
export class JupyterNotebookDiffStrategy implements DiffStrategy {
|
||||
private fallbackStrategy: MultiSearchReplaceDiffStrategy
|
||||
private globalState?: GlobalState
|
||||
|
||||
constructor(fuzzyThreshold?: number, bufferLines?: number) {
|
||||
constructor(fuzzyThreshold?: number, bufferLines?: number, globalState?: GlobalState) {
|
||||
// Use MultiSearchReplaceDiffStrategy as fallback for non-cell operations
|
||||
this.fallbackStrategy = new MultiSearchReplaceDiffStrategy(fuzzyThreshold, bufferLines)
|
||||
this.globalState = globalState
|
||||
}
|
||||
|
||||
getName(): string {
|
||||
|
|
@ -122,12 +125,28 @@ Your cell operation or search/replace content here
|
|||
// Check if this is a Jupyter notebook by trying to parse it
|
||||
let handler: JupyterNotebookHandler
|
||||
try {
|
||||
handler = new JupyterNotebookHandler("", originalContent)
|
||||
// Create security config based on global settings
|
||||
const securityConfig: JupyterSecurityConfig = {
|
||||
yoloMode: this.globalState?.jupyterNotebookYoloMode === true,
|
||||
allowCodeExecution: this.globalState?.jupyterNotebookYoloMode === true,
|
||||
readOnlyMode: this.globalState?.jupyterNotebookYoloMode !== true,
|
||||
}
|
||||
handler = new JupyterNotebookHandler("", originalContent, securityConfig)
|
||||
} catch (error) {
|
||||
// Not a valid notebook, fall back to standard diff
|
||||
return this.fallbackStrategy.applyDiff(originalContent, diffContent, _paramStartLine, _paramEndLine)
|
||||
}
|
||||
|
||||
// Check if notebook is in read-only mode due to security
|
||||
if (handler.isReadOnly()) {
|
||||
const risks = handler.getSecurityRisks()
|
||||
const riskSummary = risks.map((r) => `${r.severity}: ${r.description}`).join(", ")
|
||||
return {
|
||||
success: false,
|
||||
error: `Cannot modify notebook: Security risks detected (${riskSummary}). Enable YOLO Mode in settings to bypass security restrictions.`,
|
||||
}
|
||||
}
|
||||
|
||||
// Check if this is a cell operation
|
||||
const cellOperationMatch = diffContent.match(
|
||||
/<<<<<<< CELL_OPERATION\s*\n(?::operation:\s*(edit|add|delete)\s*\n)?(?::cell_index:\s*(\d+)\s*\n)?(?::cell_type:\s*(code|markdown|raw)\s*\n)?(?:-------\s*\n)?([\s\S]*?)(?:\n)?=======\s*\n([\s\S]*?)(?:\n)?>>>>>>> CELL_OPERATION/,
|
||||
|
|
|
|||
|
|
@ -2705,17 +2705,23 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
|
|||
const workspaceDir = vscode.workspace.workspaceFolders?.[0]?.uri.fsPath
|
||||
const hasJupyterFiles = workspaceDir && (await this.checkForJupyterFiles(workspaceDir))
|
||||
|
||||
// Get the provider and state for accessing global settings
|
||||
const provider = this.providerRef.deref()
|
||||
const state = provider ? await provider.getState() : undefined
|
||||
|
||||
if (hasJupyterFiles) {
|
||||
// Use Jupyter-specific diff strategy for notebooks
|
||||
this.diffStrategy = new JupyterNotebookDiffStrategy(this.fuzzyMatchThreshold)
|
||||
this.diffStrategy = new JupyterNotebookDiffStrategy(
|
||||
this.fuzzyMatchThreshold,
|
||||
undefined, // bufferLines
|
||||
state, // Pass global state for security settings
|
||||
)
|
||||
} else {
|
||||
// Default to old strategy, will be updated if experiment is enabled.
|
||||
this.diffStrategy = new MultiSearchReplaceDiffStrategy(this.fuzzyMatchThreshold)
|
||||
|
||||
// Check if the multi-file apply diff experiment is enabled
|
||||
const provider = this.providerRef.deref()
|
||||
if (provider) {
|
||||
const state = await provider.getState()
|
||||
if (provider && state) {
|
||||
const isMultiFileApplyDiffEnabled = experiments.isEnabled(
|
||||
state.experiments ?? {},
|
||||
EXPERIMENT_IDS.MULTI_FILE_APPLY_DIFF,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,617 @@
|
|||
import { describe, it, expect, beforeEach } from "vitest"
|
||||
import { JupyterNotebookSecurity, JupyterSecurityConfig } from "../jupyter-notebook-security"
|
||||
import { JupyterNotebook } from "../jupyter-notebook-handler"
|
||||
|
||||
describe("JupyterNotebookSecurity", () => {
|
||||
let security: JupyterNotebookSecurity
|
||||
let sampleNotebook: JupyterNotebook
|
||||
|
||||
beforeEach(() => {
|
||||
security = new JupyterNotebookSecurity()
|
||||
sampleNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "code",
|
||||
source: "print('Hello, World!')",
|
||||
metadata: {},
|
||||
outputs: [],
|
||||
execution_count: 1,
|
||||
},
|
||||
{
|
||||
cell_type: "markdown",
|
||||
source: "# Safe Markdown\nThis is safe content.",
|
||||
metadata: {},
|
||||
},
|
||||
],
|
||||
metadata: {
|
||||
kernelspec: {
|
||||
display_name: "Python 3",
|
||||
language: "python",
|
||||
name: "python3",
|
||||
},
|
||||
},
|
||||
nbformat: 4,
|
||||
nbformat_minor: 5,
|
||||
}
|
||||
})
|
||||
|
||||
describe("Safe notebooks", () => {
|
||||
it("should validate safe notebook as secure", () => {
|
||||
const result = security.validateNotebook(sampleNotebook)
|
||||
expect(result.isSecure).toBe(true)
|
||||
expect(result.risks).toHaveLength(0)
|
||||
expect(result.requiresReadOnly).toBe(false)
|
||||
})
|
||||
|
||||
it("should allow basic Python operations", () => {
|
||||
const notebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "code",
|
||||
source: "import numpy as np\nimport pandas as pd\ndata = [1, 2, 3]\nprint(sum(data))",
|
||||
metadata: {},
|
||||
},
|
||||
],
|
||||
}
|
||||
const result = security.validateNotebook(notebook)
|
||||
expect(result.isSecure).toBe(true)
|
||||
expect(result.risks).toHaveLength(0)
|
||||
})
|
||||
})
|
||||
|
||||
describe("Dangerous code patterns", () => {
|
||||
it("should detect eval() usage", () => {
|
||||
const notebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "code",
|
||||
source: "result = eval('2 + 2')",
|
||||
metadata: {},
|
||||
},
|
||||
],
|
||||
}
|
||||
const result = security.validateNotebook(notebook)
|
||||
expect(result.isSecure).toBe(false)
|
||||
expect(result.requiresReadOnly).toBe(true)
|
||||
expect(result.risks).toContainEqual(
|
||||
expect.objectContaining({
|
||||
severity: "critical",
|
||||
type: "eval",
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should detect exec() usage", () => {
|
||||
const notebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "code",
|
||||
source: "exec('import os')",
|
||||
metadata: {},
|
||||
},
|
||||
],
|
||||
}
|
||||
const result = security.validateNotebook(notebook)
|
||||
expect(result.isSecure).toBe(false)
|
||||
expect(result.requiresReadOnly).toBe(true)
|
||||
expect(result.risks).toContainEqual(
|
||||
expect.objectContaining({
|
||||
severity: "critical",
|
||||
type: "exec",
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should detect shell commands", () => {
|
||||
const notebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "code",
|
||||
source: "!rm -rf /",
|
||||
metadata: {},
|
||||
},
|
||||
],
|
||||
}
|
||||
const result = security.validateNotebook(notebook)
|
||||
expect(result.isSecure).toBe(false)
|
||||
expect(result.requiresReadOnly).toBe(true)
|
||||
expect(result.risks).toContainEqual(
|
||||
expect.objectContaining({
|
||||
severity: "critical",
|
||||
type: "shell_command",
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should detect dangerous imports", () => {
|
||||
const notebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "code",
|
||||
source: "import subprocess\nsubprocess.run(['ls', '-la'])",
|
||||
metadata: {},
|
||||
},
|
||||
],
|
||||
}
|
||||
const result = security.validateNotebook(notebook)
|
||||
expect(result.isSecure).toBe(false)
|
||||
expect(result.requiresReadOnly).toBe(true)
|
||||
expect(result.risks).toContainEqual(
|
||||
expect.objectContaining({
|
||||
severity: "high",
|
||||
type: "subprocess_import",
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should detect file operations", () => {
|
||||
const notebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "code",
|
||||
source: "with open('/etc/passwd', 'r') as f:\n content = f.read()",
|
||||
metadata: {},
|
||||
},
|
||||
],
|
||||
}
|
||||
const result = security.validateNotebook(notebook)
|
||||
expect(result.isSecure).toBe(true) // Medium risk doesn't make it insecure
|
||||
expect(result.requiresReadOnly).toBe(false)
|
||||
expect(result.risks).toContainEqual(
|
||||
expect.objectContaining({
|
||||
severity: "medium",
|
||||
type: "file_open",
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should detect network operations", () => {
|
||||
const notebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "code",
|
||||
source: "import requests\nresponse = requests.get('http://example.com')",
|
||||
metadata: {},
|
||||
},
|
||||
],
|
||||
}
|
||||
const result = security.validateNotebook(notebook)
|
||||
expect(result.isSecure).toBe(true) // Medium risk doesn't make it insecure
|
||||
expect(result.requiresReadOnly).toBe(false)
|
||||
expect(result.risks).toContainEqual(
|
||||
expect.objectContaining({
|
||||
severity: "medium",
|
||||
type: "network_request",
|
||||
}),
|
||||
)
|
||||
})
|
||||
})
|
||||
|
||||
describe("Dangerous markdown patterns", () => {
|
||||
it("should detect script tags in markdown", () => {
|
||||
const notebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "markdown",
|
||||
source: "# Title\n<script>alert('XSS')</script>",
|
||||
metadata: {},
|
||||
},
|
||||
],
|
||||
}
|
||||
const result = security.validateNotebook(notebook)
|
||||
expect(result.isSecure).toBe(false)
|
||||
expect(result.requiresReadOnly).toBe(true)
|
||||
expect(result.risks).toContainEqual(
|
||||
expect.objectContaining({
|
||||
severity: "high",
|
||||
type: "script_tag",
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should detect iframe tags", () => {
|
||||
const notebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "markdown",
|
||||
source: "<iframe src='http://malicious.com'></iframe>",
|
||||
metadata: {},
|
||||
},
|
||||
],
|
||||
}
|
||||
const result = security.validateNotebook(notebook)
|
||||
expect(result.isSecure).toBe(false)
|
||||
expect(result.requiresReadOnly).toBe(true)
|
||||
expect(result.risks).toContainEqual(
|
||||
expect.objectContaining({
|
||||
severity: "high",
|
||||
type: "iframe_tag",
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should detect javascript: protocol", () => {
|
||||
const notebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "markdown",
|
||||
source: "[Click me](javascript:alert('XSS'))",
|
||||
metadata: {},
|
||||
},
|
||||
],
|
||||
}
|
||||
const result = security.validateNotebook(notebook)
|
||||
expect(result.isSecure).toBe(false)
|
||||
expect(result.requiresReadOnly).toBe(true)
|
||||
expect(result.risks).toContainEqual(
|
||||
expect.objectContaining({
|
||||
severity: "high",
|
||||
type: "javascript_protocol",
|
||||
}),
|
||||
)
|
||||
})
|
||||
})
|
||||
|
||||
describe("Output validation", () => {
|
||||
it("should detect dangerous HTML outputs", () => {
|
||||
const notebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "code",
|
||||
source: "print('safe')",
|
||||
metadata: {},
|
||||
outputs: [
|
||||
{
|
||||
data: {
|
||||
"text/html": "<script>alert('XSS')</script>",
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
const result = security.validateNotebook(notebook)
|
||||
expect(result.isSecure).toBe(false)
|
||||
expect(result.requiresReadOnly).toBe(true)
|
||||
expect(result.risks).toContainEqual(
|
||||
expect.objectContaining({
|
||||
severity: "high",
|
||||
type: "output_script",
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should detect JavaScript outputs", () => {
|
||||
const notebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "code",
|
||||
source: "print('safe')",
|
||||
metadata: {},
|
||||
outputs: [
|
||||
{
|
||||
data: {
|
||||
"application/javascript": "console.log('executed')",
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
const result = security.validateNotebook(notebook)
|
||||
expect(result.isSecure).toBe(false)
|
||||
expect(result.requiresReadOnly).toBe(true)
|
||||
expect(result.risks).toContainEqual(
|
||||
expect.objectContaining({
|
||||
severity: "high",
|
||||
type: "javascript_output",
|
||||
}),
|
||||
)
|
||||
})
|
||||
})
|
||||
|
||||
describe("YOLO Mode", () => {
|
||||
it("should bypass all security checks when YOLO Mode is enabled", () => {
|
||||
const yoloSecurity = new JupyterNotebookSecurity({ yoloMode: true })
|
||||
const dangerousNotebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "code",
|
||||
source: 'eval(\'__import__("os").system("rm -rf /")\')',
|
||||
metadata: {},
|
||||
},
|
||||
{
|
||||
cell_type: "markdown",
|
||||
source: "<script>alert('XSS')</script>",
|
||||
metadata: {},
|
||||
},
|
||||
],
|
||||
}
|
||||
const result = yoloSecurity.validateNotebook(dangerousNotebook)
|
||||
expect(result.isSecure).toBe(true)
|
||||
expect(result.risks).toHaveLength(0)
|
||||
expect(result.requiresReadOnly).toBe(false)
|
||||
})
|
||||
|
||||
it("should allow updating YOLO Mode configuration", () => {
|
||||
const notebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "code",
|
||||
source: "eval('2+2')",
|
||||
metadata: {},
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
// Initially, security is enforced
|
||||
let result = security.validateNotebook(notebook)
|
||||
expect(result.isSecure).toBe(false)
|
||||
expect(result.requiresReadOnly).toBe(true)
|
||||
|
||||
// Enable YOLO Mode
|
||||
security.updateConfig({ yoloMode: true })
|
||||
result = security.validateNotebook(notebook)
|
||||
expect(result.isSecure).toBe(true)
|
||||
expect(result.requiresReadOnly).toBe(false)
|
||||
|
||||
// Disable YOLO Mode
|
||||
security.updateConfig({ yoloMode: false })
|
||||
result = security.validateNotebook(notebook)
|
||||
expect(result.isSecure).toBe(false)
|
||||
expect(result.requiresReadOnly).toBe(true)
|
||||
})
|
||||
|
||||
it("should correctly report YOLO Mode status", () => {
|
||||
expect(security.isYoloModeEnabled()).toBe(false)
|
||||
|
||||
security.updateConfig({ yoloMode: true })
|
||||
expect(security.isYoloModeEnabled()).toBe(true)
|
||||
|
||||
security.updateConfig({ yoloMode: false })
|
||||
expect(security.isYoloModeEnabled()).toBe(false)
|
||||
})
|
||||
})
|
||||
|
||||
describe("Sanitization", () => {
|
||||
it("should sanitize dangerous code cells", () => {
|
||||
const notebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "code",
|
||||
source: "eval('malicious code')",
|
||||
metadata: {},
|
||||
outputs: [{ data: { "text/plain": "output" } }],
|
||||
execution_count: 1,
|
||||
},
|
||||
],
|
||||
}
|
||||
const result = security.validateNotebook(notebook)
|
||||
expect(result.sanitizedNotebook).toBeDefined()
|
||||
const sanitizedCell = result.sanitizedNotebook!.cells[0]
|
||||
expect(sanitizedCell.source).toContain("SECURITY WARNING")
|
||||
expect(sanitizedCell.source).toContain("eval")
|
||||
expect(sanitizedCell.outputs).toHaveLength(0)
|
||||
expect(sanitizedCell.execution_count).toBeNull()
|
||||
})
|
||||
|
||||
it("should sanitize dangerous markdown cells", () => {
|
||||
const notebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "markdown",
|
||||
source: "# Title\n<script>alert('XSS')</script>\n<iframe src='bad'></iframe>",
|
||||
metadata: {},
|
||||
},
|
||||
],
|
||||
}
|
||||
const result = security.validateNotebook(notebook)
|
||||
expect(result.sanitizedNotebook).toBeDefined()
|
||||
const sanitizedCell = result.sanitizedNotebook!.cells[0]
|
||||
expect(sanitizedCell.source).toContain("[REMOVED: script tag]")
|
||||
expect(sanitizedCell.source).toContain("[REMOVED: iframe]")
|
||||
expect(sanitizedCell.source).not.toContain("<script>")
|
||||
expect(sanitizedCell.source).not.toContain("<iframe")
|
||||
})
|
||||
|
||||
it("should remove dangerous outputs", () => {
|
||||
const notebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "code",
|
||||
source: "import os", // High risk to trigger sanitization
|
||||
metadata: {},
|
||||
outputs: [
|
||||
{
|
||||
data: {
|
||||
"text/html": "<script>alert('XSS')</script>",
|
||||
"application/javascript": "console.log('bad')",
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
const result = security.validateNotebook(notebook)
|
||||
expect(result.sanitizedNotebook).toBeDefined()
|
||||
|
||||
// The cell should be sanitized because it has high-risk import
|
||||
const sanitizedCell = result.sanitizedNotebook!.cells[0]
|
||||
expect(sanitizedCell.source).toContain("SECURITY WARNING")
|
||||
|
||||
// Outputs should be cleared for dangerous cells
|
||||
expect(sanitizedCell.outputs).toHaveLength(0)
|
||||
})
|
||||
|
||||
it("should sanitize dangerous outputs in safe cells", () => {
|
||||
const notebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "code",
|
||||
source: "print('safe code')", // Safe code
|
||||
metadata: {},
|
||||
outputs: [
|
||||
{
|
||||
data: {
|
||||
"text/html": "<div>Safe HTML</div>",
|
||||
"text/plain": ["safe output"],
|
||||
},
|
||||
},
|
||||
{
|
||||
data: {
|
||||
"text/html": "<script>alert('XSS')</script>",
|
||||
"application/javascript": "console.log('bad')",
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
const result = security.validateNotebook(notebook)
|
||||
|
||||
// Since the outputs have high-risk content (script tags), it should be flagged
|
||||
expect(result.isSecure).toBe(false)
|
||||
expect(result.sanitizedNotebook).toBeDefined()
|
||||
|
||||
// The cell should be disabled because it has high-risk outputs
|
||||
const sanitizedCell = result.sanitizedNotebook!.cells[0]
|
||||
expect(sanitizedCell.source).toContain("SECURITY WARNING")
|
||||
expect(sanitizedCell.outputs).toHaveLength(0)
|
||||
})
|
||||
|
||||
it("should not sanitize when YOLO Mode is enabled", () => {
|
||||
const yoloSecurity = new JupyterNotebookSecurity({ yoloMode: true })
|
||||
const notebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "code",
|
||||
source: "eval('dangerous')",
|
||||
metadata: {},
|
||||
},
|
||||
],
|
||||
}
|
||||
const result = yoloSecurity.validateNotebook(notebook)
|
||||
expect(result.sanitizedNotebook).toBeUndefined()
|
||||
})
|
||||
})
|
||||
|
||||
describe("Trusted sources", () => {
|
||||
it("should trust notebooks from trusted sources", () => {
|
||||
const securityWithTrusted = new JupyterNotebookSecurity({
|
||||
trustedSources: ["/trusted/path", "/safe/notebooks"],
|
||||
})
|
||||
const dangerousNotebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "code",
|
||||
source: "eval('trusted code')",
|
||||
metadata: {},
|
||||
},
|
||||
],
|
||||
}
|
||||
const result = securityWithTrusted.validateNotebook(dangerousNotebook, "/trusted/path/notebook.ipynb")
|
||||
expect(result.isSecure).toBe(true)
|
||||
expect(result.risks).toHaveLength(0)
|
||||
expect(result.requiresReadOnly).toBe(false)
|
||||
})
|
||||
|
||||
it("should not trust notebooks from untrusted sources", () => {
|
||||
const securityWithTrusted = new JupyterNotebookSecurity({
|
||||
trustedSources: ["/trusted/path"],
|
||||
})
|
||||
const dangerousNotebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "code",
|
||||
source: "eval('untrusted code')",
|
||||
metadata: {},
|
||||
},
|
||||
],
|
||||
}
|
||||
const result = securityWithTrusted.validateNotebook(dangerousNotebook, "/untrusted/path/notebook.ipynb")
|
||||
expect(result.isSecure).toBe(false)
|
||||
expect(result.risks.length).toBeGreaterThan(0)
|
||||
expect(result.requiresReadOnly).toBe(true)
|
||||
})
|
||||
})
|
||||
|
||||
describe("Size limits", () => {
|
||||
it("should detect oversized cells", () => {
|
||||
const largeContent = "x".repeat(2 * 1024 * 1024) // 2MB
|
||||
const notebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "code",
|
||||
source: largeContent,
|
||||
metadata: {},
|
||||
},
|
||||
],
|
||||
}
|
||||
const result = security.validateNotebook(notebook)
|
||||
expect(result.risks).toContainEqual(
|
||||
expect.objectContaining({
|
||||
severity: "medium",
|
||||
type: "oversized_cell",
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should detect excessive cell count", () => {
|
||||
const cells = Array(1001).fill({
|
||||
cell_type: "code",
|
||||
source: "print('cell')",
|
||||
metadata: {},
|
||||
})
|
||||
const notebook: JupyterNotebook = { cells }
|
||||
const result = security.validateNotebook(notebook)
|
||||
expect(result.risks).toContainEqual(
|
||||
expect.objectContaining({
|
||||
severity: "medium",
|
||||
type: "excessive_cells",
|
||||
}),
|
||||
)
|
||||
})
|
||||
})
|
||||
|
||||
describe("Metadata validation", () => {
|
||||
it("should detect suspicious metadata fields", () => {
|
||||
const notebook: JupyterNotebook = {
|
||||
cells: [],
|
||||
metadata: {
|
||||
widgets: { some: "data" },
|
||||
extensions: { another: "data" },
|
||||
jupyter_dashboards: { layout: "grid" },
|
||||
},
|
||||
}
|
||||
const result = security.validateNotebook(notebook)
|
||||
expect(result.risks).toHaveLength(3)
|
||||
expect(result.risks).toContainEqual(
|
||||
expect.objectContaining({
|
||||
severity: "low",
|
||||
type: "suspicious_metadata",
|
||||
description: expect.stringContaining("widgets"),
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should sanitize suspicious metadata", () => {
|
||||
const notebook: JupyterNotebook = {
|
||||
cells: [
|
||||
{
|
||||
cell_type: "code",
|
||||
source: "eval('bad')",
|
||||
metadata: {},
|
||||
},
|
||||
],
|
||||
metadata: {
|
||||
widgets: { some: "data" },
|
||||
kernelspec: { name: "python3" },
|
||||
},
|
||||
}
|
||||
const result = security.validateNotebook(notebook)
|
||||
expect(result.sanitizedNotebook).toBeDefined()
|
||||
expect(result.sanitizedNotebook!.metadata?.widgets).toBeUndefined()
|
||||
expect(result.sanitizedNotebook!.metadata?.kernelspec).toBeDefined()
|
||||
})
|
||||
})
|
||||
})
|
||||
|
|
@ -8,6 +8,7 @@ import { extractTextFromXLSX } from "./extract-text-from-xlsx"
|
|||
import { countFileLines } from "./line-counter"
|
||||
import { readLines } from "./read-lines"
|
||||
import { JupyterNotebookHandler } from "./jupyter-notebook-handler"
|
||||
import { JupyterSecurityConfig } from "./jupyter-notebook-security"
|
||||
|
||||
async function extractTextFromPDF(filePath: string): Promise<string> {
|
||||
const dataBuffer = await fs.readFile(filePath)
|
||||
|
|
@ -20,9 +21,32 @@ async function extractTextFromDOCX(filePath: string): Promise<string> {
|
|||
return addLineNumbers(result.value)
|
||||
}
|
||||
|
||||
async function extractTextFromIPYNB(filePath: string): Promise<string> {
|
||||
const handler = await JupyterNotebookHandler.fromFile(filePath)
|
||||
return handler.extractTextWithCellMarkers()
|
||||
async function extractTextFromIPYNB(filePath: string, securityConfig?: JupyterSecurityConfig): Promise<string> {
|
||||
const handler = await JupyterNotebookHandler.fromFile(filePath, securityConfig)
|
||||
|
||||
// Add security warning if risks are detected
|
||||
let result = ""
|
||||
if (handler.hasSecurityRisks() && !securityConfig?.yoloMode) {
|
||||
const risks = handler.getSecurityRisks()
|
||||
const criticalRisks = risks.filter((r) => r.severity === "critical")
|
||||
const highRisks = risks.filter((r) => r.severity === "high")
|
||||
|
||||
result += "# ⚠️ SECURITY WARNING ⚠️\n"
|
||||
result += "# This Jupyter notebook contains potential security risks:\n"
|
||||
|
||||
if (criticalRisks.length > 0) {
|
||||
result += `# - ${criticalRisks.length} CRITICAL risk(s): ${criticalRisks.map((r) => r.type).join(", ")}\n`
|
||||
}
|
||||
if (highRisks.length > 0) {
|
||||
result += `# - ${highRisks.length} HIGH risk(s): ${highRisks.map((r) => r.type).join(", ")}\n`
|
||||
}
|
||||
|
||||
result += "# The notebook is in READ-ONLY mode. Enable YOLO Mode to bypass restrictions.\n"
|
||||
result += "# " + "=".repeat(70) + "\n\n"
|
||||
}
|
||||
|
||||
result += handler.extractTextWithCellMarkers()
|
||||
return result
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import * as fs from "fs/promises"
|
||||
import * as path from "path"
|
||||
import { addLineNumbers } from "./extract-text"
|
||||
import { JupyterNotebookSecurity, SecurityValidationResult, JupyterSecurityConfig } from "./jupyter-notebook-security"
|
||||
|
||||
export interface JupyterCell {
|
||||
cell_type: "code" | "markdown" | "raw"
|
||||
|
|
@ -29,12 +30,23 @@ export class JupyterNotebookHandler {
|
|||
private notebook: JupyterNotebook
|
||||
private filePath: string
|
||||
private cellReferences: CellReference[] = []
|
||||
private security: JupyterNotebookSecurity
|
||||
private securityValidation?: SecurityValidationResult
|
||||
|
||||
constructor(filePath: string, notebookContent?: string) {
|
||||
constructor(filePath: string, notebookContent?: string, securityConfig?: JupyterSecurityConfig) {
|
||||
this.filePath = filePath
|
||||
this.security = new JupyterNotebookSecurity(securityConfig)
|
||||
|
||||
if (notebookContent) {
|
||||
this.notebook = JSON.parse(notebookContent)
|
||||
this.buildCellReferences()
|
||||
// Validate security on load
|
||||
this.securityValidation = this.security.validateNotebook(this.notebook, filePath)
|
||||
// Use sanitized notebook if available and not in YOLO mode
|
||||
if (this.securityValidation.sanitizedNotebook && !this.security.isYoloModeEnabled()) {
|
||||
this.notebook = this.securityValidation.sanitizedNotebook
|
||||
this.buildCellReferences()
|
||||
}
|
||||
} else {
|
||||
this.notebook = { cells: [] }
|
||||
}
|
||||
|
|
@ -43,9 +55,9 @@ export class JupyterNotebookHandler {
|
|||
/**
|
||||
* Load a Jupyter notebook from file
|
||||
*/
|
||||
static async fromFile(filePath: string): Promise<JupyterNotebookHandler> {
|
||||
static async fromFile(filePath: string, securityConfig?: JupyterSecurityConfig): Promise<JupyterNotebookHandler> {
|
||||
const content = await fs.readFile(filePath, "utf8")
|
||||
return new JupyterNotebookHandler(filePath, content)
|
||||
return new JupyterNotebookHandler(filePath, content, securityConfig)
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -140,6 +152,12 @@ export class JupyterNotebookHandler {
|
|||
* Update a specific cell's content
|
||||
*/
|
||||
updateCell(cellIndex: number, newContent: string): boolean {
|
||||
// Check if read-only mode is enforced
|
||||
if (this.isReadOnly()) {
|
||||
console.warn("Cannot update cell: Notebook is in read-only mode due to security restrictions")
|
||||
return false
|
||||
}
|
||||
|
||||
if (cellIndex < 0 || cellIndex >= this.notebook.cells.length) {
|
||||
return false
|
||||
}
|
||||
|
|
@ -167,6 +185,12 @@ export class JupyterNotebookHandler {
|
|||
* Insert a new cell
|
||||
*/
|
||||
insertCell(index: number, cellType: "code" | "markdown" | "raw", content: string): boolean {
|
||||
// Check if read-only mode is enforced
|
||||
if (this.isReadOnly()) {
|
||||
console.warn("Cannot insert cell: Notebook is in read-only mode due to security restrictions")
|
||||
return false
|
||||
}
|
||||
|
||||
if (index < 0 || index > this.notebook.cells.length) {
|
||||
return false
|
||||
}
|
||||
|
|
@ -193,6 +217,12 @@ export class JupyterNotebookHandler {
|
|||
* Delete a cell
|
||||
*/
|
||||
deleteCell(index: number): boolean {
|
||||
// Check if read-only mode is enforced
|
||||
if (this.isReadOnly()) {
|
||||
console.warn("Cannot delete cell: Notebook is in read-only mode due to security restrictions")
|
||||
return false
|
||||
}
|
||||
|
||||
if (index < 0 || index >= this.notebook.cells.length) {
|
||||
return false
|
||||
}
|
||||
|
|
@ -226,6 +256,11 @@ export class JupyterNotebookHandler {
|
|||
* Save the notebook back to file
|
||||
*/
|
||||
async save(): Promise<void> {
|
||||
// Check if read-only mode is enforced
|
||||
if (this.isReadOnly()) {
|
||||
throw new Error("Cannot save notebook: Notebook is in read-only mode due to security restrictions")
|
||||
}
|
||||
|
||||
const content = JSON.stringify(this.notebook, null, 2)
|
||||
await fs.writeFile(this.filePath, content, "utf8")
|
||||
}
|
||||
|
|
@ -335,4 +370,41 @@ export class JupyterNotebookHandler {
|
|||
cells,
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get security validation results
|
||||
*/
|
||||
getSecurityValidation(): SecurityValidationResult | undefined {
|
||||
return this.securityValidation
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if notebook is in read-only mode
|
||||
*/
|
||||
isReadOnly(): boolean {
|
||||
return this.securityValidation?.requiresReadOnly === true && !this.security.isYoloModeEnabled()
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if notebook has security risks
|
||||
*/
|
||||
hasSecurityRisks(): boolean {
|
||||
return (this.securityValidation?.risks.length ?? 0) > 0
|
||||
}
|
||||
|
||||
/**
|
||||
* Get security risks
|
||||
*/
|
||||
getSecurityRisks() {
|
||||
return this.securityValidation?.risks || []
|
||||
}
|
||||
|
||||
/**
|
||||
* Enable or disable YOLO Mode
|
||||
*/
|
||||
setYoloMode(enabled: boolean): void {
|
||||
this.security.updateConfig({ yoloMode: enabled })
|
||||
// Re-validate with new settings
|
||||
this.securityValidation = this.security.validateNotebook(this.notebook, this.filePath)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
418
src/integrations/misc/jupyter-notebook-security.ts
Normal file
418
src/integrations/misc/jupyter-notebook-security.ts
Normal file
|
|
@ -0,0 +1,418 @@
|
|||
import { JupyterNotebook, JupyterCell } from "./jupyter-notebook-handler"
|
||||
|
||||
export interface SecurityRisk {
|
||||
severity: "critical" | "high" | "medium" | "low"
|
||||
type: string
|
||||
cellIndex?: number
|
||||
cellType?: string
|
||||
description: string
|
||||
pattern?: string
|
||||
}
|
||||
|
||||
export interface SecurityValidationResult {
|
||||
isSecure: boolean
|
||||
risks: SecurityRisk[]
|
||||
requiresReadOnly: boolean
|
||||
sanitizedNotebook?: JupyterNotebook
|
||||
}
|
||||
|
||||
export interface JupyterSecurityConfig {
|
||||
allowCodeExecution?: boolean
|
||||
readOnlyMode?: boolean
|
||||
maxCellSize?: number
|
||||
maxCellCount?: number
|
||||
trustedSources?: string[]
|
||||
yoloMode?: boolean // New YOLO Mode flag
|
||||
}
|
||||
|
||||
const DEFAULT_CONFIG: JupyterSecurityConfig = {
|
||||
allowCodeExecution: false,
|
||||
readOnlyMode: false, // Only enforce read-only when risks are detected
|
||||
maxCellSize: 1024 * 1024, // 1MB
|
||||
maxCellCount: 1000,
|
||||
trustedSources: [],
|
||||
yoloMode: false,
|
||||
}
|
||||
|
||||
// Dangerous code patterns that could execute arbitrary code
|
||||
const DANGEROUS_CODE_PATTERNS = [
|
||||
// Direct code execution
|
||||
{ pattern: /\beval\s*\(/, severity: "critical" as const, type: "eval" },
|
||||
{ pattern: /\bexec\s*\(/, severity: "critical" as const, type: "exec" },
|
||||
{ pattern: /\bcompile\s*\(/, severity: "critical" as const, type: "compile" },
|
||||
{ pattern: /\b__import__\s*\(/, severity: "critical" as const, type: "__import__" },
|
||||
|
||||
// System commands
|
||||
{ pattern: /^!.*/, severity: "critical" as const, type: "shell_command" },
|
||||
{ pattern: /%system\s+/, severity: "critical" as const, type: "magic_system" },
|
||||
{ pattern: /%%bash/, severity: "critical" as const, type: "magic_bash" },
|
||||
{ pattern: /%%sh/, severity: "critical" as const, type: "magic_sh" },
|
||||
{ pattern: /%%script/, severity: "critical" as const, type: "magic_script" },
|
||||
|
||||
// Dangerous imports
|
||||
{ pattern: /import\s+subprocess/, severity: "high" as const, type: "subprocess_import" },
|
||||
{ pattern: /from\s+subprocess\s+import/, severity: "high" as const, type: "subprocess_import" },
|
||||
{ pattern: /import\s+os/, severity: "high" as const, type: "os_import" },
|
||||
{ pattern: /from\s+os\s+import/, severity: "high" as const, type: "os_import" },
|
||||
{ pattern: /import\s+sys/, severity: "high" as const, type: "sys_import" },
|
||||
{ pattern: /from\s+sys\s+import/, severity: "high" as const, type: "sys_import" },
|
||||
{ pattern: /import\s+socket/, severity: "high" as const, type: "socket_import" },
|
||||
{ pattern: /from\s+socket\s+import/, severity: "high" as const, type: "socket_import" },
|
||||
{ pattern: /import\s+pickle/, severity: "high" as const, type: "pickle_import" },
|
||||
{ pattern: /from\s+pickle\s+import/, severity: "high" as const, type: "pickle_import" },
|
||||
|
||||
// File operations
|
||||
{ pattern: /\bopen\s*\(/, severity: "medium" as const, type: "file_open" },
|
||||
{ pattern: /\bfile\s*\(/, severity: "medium" as const, type: "file_operation" },
|
||||
{ pattern: /\.write\s*\(/, severity: "medium" as const, type: "file_write" },
|
||||
{ pattern: /\.read\s*\(/, severity: "medium" as const, type: "file_read" },
|
||||
|
||||
// Network operations
|
||||
{ pattern: /requests\.(get|post|put|delete|patch)/, severity: "medium" as const, type: "network_request" },
|
||||
{ pattern: /urllib\.request/, severity: "medium" as const, type: "network_urllib" },
|
||||
{ pattern: /http\.client/, severity: "medium" as const, type: "network_http" },
|
||||
]
|
||||
|
||||
// Dangerous patterns in markdown cells (potential XSS)
|
||||
const DANGEROUS_MARKDOWN_PATTERNS = [
|
||||
{ pattern: /<script[^>]*>[\s\S]*?<\/script>/gi, severity: "high" as const, type: "script_tag" },
|
||||
{ pattern: /<iframe[^>]*>/gi, severity: "high" as const, type: "iframe_tag" },
|
||||
{ pattern: /javascript:/gi, severity: "high" as const, type: "javascript_protocol" },
|
||||
{ pattern: /on\w+\s*=/gi, severity: "medium" as const, type: "event_handler" },
|
||||
]
|
||||
|
||||
// Dangerous output patterns
|
||||
const DANGEROUS_OUTPUT_PATTERNS = [
|
||||
{ pattern: /<script[^>]*>[\s\S]*?<\/script>/gi, severity: "high" as const, type: "output_script" },
|
||||
{ pattern: /data:text\/html/gi, severity: "medium" as const, type: "html_data_uri" },
|
||||
]
|
||||
|
||||
export class JupyterNotebookSecurity {
|
||||
private config: JupyterSecurityConfig
|
||||
|
||||
constructor(config?: JupyterSecurityConfig) {
|
||||
this.config = { ...DEFAULT_CONFIG, ...config }
|
||||
}
|
||||
|
||||
/**
|
||||
* Validate a Jupyter notebook for security risks
|
||||
*/
|
||||
validateNotebook(notebook: JupyterNotebook, filePath?: string): SecurityValidationResult {
|
||||
// If YOLO Mode is enabled, bypass all security checks
|
||||
if (this.config.yoloMode) {
|
||||
return {
|
||||
isSecure: true,
|
||||
risks: [],
|
||||
requiresReadOnly: false,
|
||||
}
|
||||
}
|
||||
|
||||
const risks: SecurityRisk[] = []
|
||||
|
||||
// Check if source is trusted
|
||||
if (filePath && this.config.trustedSources?.length) {
|
||||
const isTrusted = this.config.trustedSources.some((source) => filePath.includes(source))
|
||||
if (isTrusted) {
|
||||
return {
|
||||
isSecure: true,
|
||||
risks: [],
|
||||
requiresReadOnly: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check cell count
|
||||
if (notebook.cells.length > (this.config.maxCellCount || DEFAULT_CONFIG.maxCellCount!)) {
|
||||
risks.push({
|
||||
severity: "medium",
|
||||
type: "excessive_cells",
|
||||
description: `Notebook has ${notebook.cells.length} cells, exceeding limit of ${this.config.maxCellCount}`,
|
||||
})
|
||||
}
|
||||
|
||||
// Validate each cell
|
||||
notebook.cells.forEach((cell, index) => {
|
||||
const cellRisks = this.validateCell(cell, index)
|
||||
risks.push(...cellRisks)
|
||||
})
|
||||
|
||||
// Check metadata for suspicious fields
|
||||
if (notebook.metadata) {
|
||||
const metadataRisks = this.validateMetadata(notebook.metadata)
|
||||
risks.push(...metadataRisks)
|
||||
}
|
||||
|
||||
// Determine security status
|
||||
const hasCriticalRisk = risks.some((r) => r.severity === "critical")
|
||||
const hasHighRisk = risks.some((r) => r.severity === "high")
|
||||
|
||||
const isSecure = !hasCriticalRisk && !hasHighRisk
|
||||
// Only require read-only if explicitly configured OR if there are critical/high risks
|
||||
const requiresReadOnly = this.config.readOnlyMode === true || hasCriticalRisk || hasHighRisk
|
||||
|
||||
// Optionally sanitize the notebook
|
||||
let sanitizedNotebook: JupyterNotebook | undefined
|
||||
if (!isSecure && this.config.allowCodeExecution === false) {
|
||||
sanitizedNotebook = this.sanitizeNotebook(notebook, risks)
|
||||
}
|
||||
|
||||
return {
|
||||
isSecure,
|
||||
risks,
|
||||
requiresReadOnly,
|
||||
sanitizedNotebook,
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Validate a single cell for security risks
|
||||
*/
|
||||
private validateCell(cell: JupyterCell, index: number): SecurityRisk[] {
|
||||
const risks: SecurityRisk[] = []
|
||||
const source = Array.isArray(cell.source) ? cell.source.join("") : cell.source || ""
|
||||
|
||||
// Check cell size
|
||||
const cellSize = new TextEncoder().encode(source).length
|
||||
if (cellSize > (this.config.maxCellSize || DEFAULT_CONFIG.maxCellSize!)) {
|
||||
risks.push({
|
||||
severity: "medium",
|
||||
type: "oversized_cell",
|
||||
cellIndex: index,
|
||||
cellType: cell.cell_type,
|
||||
description: `Cell ${index} exceeds size limit (${cellSize} bytes)`,
|
||||
})
|
||||
}
|
||||
|
||||
// Check for dangerous patterns based on cell type
|
||||
if (cell.cell_type === "code") {
|
||||
// Check code patterns
|
||||
for (const { pattern, severity, type } of DANGEROUS_CODE_PATTERNS) {
|
||||
if (pattern.test(source)) {
|
||||
risks.push({
|
||||
severity,
|
||||
type,
|
||||
cellIndex: index,
|
||||
cellType: "code",
|
||||
description: `Dangerous code pattern detected: ${type}`,
|
||||
pattern: pattern.toString(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Check outputs for dangerous content
|
||||
if (cell.outputs && Array.isArray(cell.outputs)) {
|
||||
for (const output of cell.outputs) {
|
||||
if (output.data) {
|
||||
const outputRisks = this.validateOutputData(output.data, index)
|
||||
risks.push(...outputRisks)
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if (cell.cell_type === "markdown") {
|
||||
// Check markdown patterns
|
||||
for (const { pattern, severity, type } of DANGEROUS_MARKDOWN_PATTERNS) {
|
||||
if (pattern.test(source)) {
|
||||
risks.push({
|
||||
severity,
|
||||
type,
|
||||
cellIndex: index,
|
||||
cellType: "markdown",
|
||||
description: `Dangerous markdown pattern detected: ${type}`,
|
||||
pattern: pattern.toString(),
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return risks
|
||||
}
|
||||
|
||||
/**
|
||||
* Validate output data for security risks
|
||||
*/
|
||||
private validateOutputData(data: any, cellIndex: number): SecurityRisk[] {
|
||||
const risks: SecurityRisk[] = []
|
||||
|
||||
// Check HTML outputs
|
||||
if (data["text/html"]) {
|
||||
const htmlContent = Array.isArray(data["text/html"]) ? data["text/html"].join("") : data["text/html"]
|
||||
|
||||
for (const { pattern, severity, type } of DANGEROUS_OUTPUT_PATTERNS) {
|
||||
if (pattern.test(htmlContent)) {
|
||||
risks.push({
|
||||
severity,
|
||||
type,
|
||||
cellIndex,
|
||||
cellType: "output",
|
||||
description: `Dangerous output pattern detected: ${type}`,
|
||||
pattern: pattern.toString(),
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check JavaScript outputs
|
||||
if (data["application/javascript"]) {
|
||||
risks.push({
|
||||
severity: "high",
|
||||
type: "javascript_output",
|
||||
cellIndex,
|
||||
cellType: "output",
|
||||
description: "JavaScript output detected",
|
||||
})
|
||||
}
|
||||
|
||||
return risks
|
||||
}
|
||||
|
||||
/**
|
||||
* Validate notebook metadata for security risks
|
||||
*/
|
||||
private validateMetadata(metadata: Record<string, any>): SecurityRisk[] {
|
||||
const risks: SecurityRisk[] = []
|
||||
|
||||
// Check for suspicious metadata fields
|
||||
const suspiciousFields = ["widgets", "extensions", "jupyter_dashboards"]
|
||||
for (const field of suspiciousFields) {
|
||||
if (metadata[field]) {
|
||||
risks.push({
|
||||
severity: "low",
|
||||
type: "suspicious_metadata",
|
||||
description: `Suspicious metadata field detected: ${field}`,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return risks
|
||||
}
|
||||
|
||||
/**
|
||||
* Sanitize a notebook by removing or disabling dangerous content
|
||||
*/
|
||||
private sanitizeNotebook(notebook: JupyterNotebook, risks: SecurityRisk[]): JupyterNotebook {
|
||||
const sanitized = JSON.parse(JSON.stringify(notebook)) as JupyterNotebook
|
||||
|
||||
// Group risks by cell index
|
||||
const risksByCell = new Map<number, SecurityRisk[]>()
|
||||
for (const risk of risks) {
|
||||
if (risk.cellIndex !== undefined) {
|
||||
if (!risksByCell.has(risk.cellIndex)) {
|
||||
risksByCell.set(risk.cellIndex, [])
|
||||
}
|
||||
risksByCell.get(risk.cellIndex)!.push(risk)
|
||||
}
|
||||
}
|
||||
|
||||
// Sanitize cells with risks
|
||||
for (const [cellIndex, cellRisks] of risksByCell) {
|
||||
const cell = sanitized.cells[cellIndex]
|
||||
if (!cell) continue
|
||||
|
||||
const hasCriticalRisk = cellRisks.some((r) => r.severity === "critical")
|
||||
const hasHighRisk = cellRisks.some((r) => r.severity === "high")
|
||||
|
||||
if (cell.cell_type === "code" && (hasCriticalRisk || hasHighRisk)) {
|
||||
// Disable dangerous code cells
|
||||
const riskTypes = cellRisks.map((r) => r.type).join(", ")
|
||||
const warningComment = `# ⚠️ SECURITY WARNING: This cell has been disabled due to security risks (${riskTypes})\n# To run this cell, enable YOLO Mode in settings\n\n`
|
||||
|
||||
if (Array.isArray(cell.source)) {
|
||||
const originalCode = cell.source.map((line) => `# ${line}`)
|
||||
cell.source = [warningComment, "# Original code:\n", ...originalCode]
|
||||
} else {
|
||||
const originalCode = (cell.source || "")
|
||||
.split("\n")
|
||||
.map((line) => `# ${line}`)
|
||||
.join("\n")
|
||||
cell.source = warningComment + "# Original code:\n" + originalCode
|
||||
}
|
||||
|
||||
// Clear outputs
|
||||
cell.outputs = []
|
||||
cell.execution_count = null
|
||||
} else if (cell.cell_type === "markdown" && (hasCriticalRisk || hasHighRisk)) {
|
||||
// Sanitize dangerous markdown
|
||||
let source = Array.isArray(cell.source) ? cell.source.join("") : cell.source || ""
|
||||
|
||||
// Remove script tags
|
||||
source = source.replace(/<script[^>]*>[\s\S]*?<\/script>/gi, "[REMOVED: script tag]")
|
||||
// Remove iframes
|
||||
source = source.replace(/<iframe[^>]*>[\s\S]*?<\/iframe>/gi, "[REMOVED: iframe]")
|
||||
// Remove javascript: protocols
|
||||
source = source.replace(/javascript:/gi, "[REMOVED]:")
|
||||
// Remove event handlers
|
||||
source = source.replace(/on\w+\s*=/gi, "data-removed=")
|
||||
|
||||
// Handle both array and string source formats
|
||||
if (Array.isArray(cell.source)) {
|
||||
cell.source = source
|
||||
.split("\n")
|
||||
.map((line, idx, arr) => (idx === arr.length - 1 && line === "" ? line : line + "\n"))
|
||||
} else {
|
||||
cell.source = source
|
||||
}
|
||||
}
|
||||
|
||||
// Clear dangerous outputs
|
||||
if (cell.outputs && Array.isArray(cell.outputs)) {
|
||||
cell.outputs = cell.outputs.map((output) => {
|
||||
if (output && output.data) {
|
||||
const newOutput = { ...output, data: { ...output.data } }
|
||||
|
||||
// Remove HTML outputs with scripts
|
||||
if (newOutput.data["text/html"]) {
|
||||
const html = Array.isArray(newOutput.data["text/html"])
|
||||
? newOutput.data["text/html"].join("")
|
||||
: newOutput.data["text/html"]
|
||||
|
||||
if (/<script|<iframe|javascript:/i.test(html)) {
|
||||
delete newOutput.data["text/html"]
|
||||
// Only add text/plain if it doesn't exist
|
||||
if (!newOutput.data["text/plain"]) {
|
||||
newOutput.data["text/plain"] = ["[HTML output removed for security]"]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Remove JavaScript outputs
|
||||
if (newOutput.data["application/javascript"]) {
|
||||
delete newOutput.data["application/javascript"]
|
||||
// Add or append to text/plain
|
||||
if (!newOutput.data["text/plain"]) {
|
||||
newOutput.data["text/plain"] = ["[JavaScript output removed for security]"]
|
||||
} else if (Array.isArray(newOutput.data["text/plain"])) {
|
||||
newOutput.data["text/plain"].push("[JavaScript output removed for security]")
|
||||
}
|
||||
}
|
||||
|
||||
return newOutput
|
||||
}
|
||||
return output
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Remove suspicious metadata
|
||||
if (sanitized.metadata) {
|
||||
delete sanitized.metadata.widgets
|
||||
delete sanitized.metadata.extensions
|
||||
delete sanitized.metadata.jupyter_dashboards
|
||||
}
|
||||
|
||||
return sanitized
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if YOLO Mode is enabled
|
||||
*/
|
||||
isYoloModeEnabled(): boolean {
|
||||
return this.config.yoloMode === true
|
||||
}
|
||||
|
||||
/**
|
||||
* Update security configuration
|
||||
*/
|
||||
updateConfig(config: Partial<JupyterSecurityConfig>): void {
|
||||
this.config = { ...this.config, ...config }
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue