diff --git a/docs/jupyter-notebook-security.md b/docs/jupyter-notebook-security.md new file mode 100644 index 0000000000..c7f3e88b1e --- /dev/null +++ b/docs/jupyter-notebook-security.md @@ -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**: ` + +""" +``` + +### 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 diff --git a/packages/types/src/global-settings.ts b/packages/types/src/global-settings.ts index 81c6ae6dfe..492e5d6dc4 100644 --- a/packages/types/src/global-settings.ts +++ b/packages/types/src/global-settings.ts @@ -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 diff --git a/src/core/diff/strategies/jupyter-notebook-diff.ts b/src/core/diff/strategies/jupyter-notebook-diff.ts index 031c230a63..059bf25c4f 100644 --- a/src/core/diff/strategies/jupyter-notebook-diff.ts +++ b/src/core/diff/strategies/jupyter-notebook-diff.ts @@ -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/, diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index 615079da75..a49422ce5e 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -2705,17 +2705,23 @@ export class Task extends EventEmitter 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, diff --git a/src/integrations/misc/__tests__/jupyter-notebook-security.spec.ts b/src/integrations/misc/__tests__/jupyter-notebook-security.spec.ts new file mode 100644 index 0000000000..678cffc3a8 --- /dev/null +++ b/src/integrations/misc/__tests__/jupyter-notebook-security.spec.ts @@ -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", + 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: "", + 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": "", + }, + }, + ], + }, + ], + } + 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: "", + 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\n", + 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("", + "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": "
Safe HTML
", + "text/plain": ["safe output"], + }, + }, + { + data: { + "text/html": "", + "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() + }) + }) +}) diff --git a/src/integrations/misc/extract-text.ts b/src/integrations/misc/extract-text.ts index 58d166153e..e10df22ef0 100644 --- a/src/integrations/misc/extract-text.ts +++ b/src/integrations/misc/extract-text.ts @@ -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 { const dataBuffer = await fs.readFile(filePath) @@ -20,9 +21,32 @@ async function extractTextFromDOCX(filePath: string): Promise { return addLineNumbers(result.value) } -async function extractTextFromIPYNB(filePath: string): Promise { - const handler = await JupyterNotebookHandler.fromFile(filePath) - return handler.extractTextWithCellMarkers() +async function extractTextFromIPYNB(filePath: string, securityConfig?: JupyterSecurityConfig): Promise { + 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 } /** diff --git a/src/integrations/misc/jupyter-notebook-handler.ts b/src/integrations/misc/jupyter-notebook-handler.ts index 529dff5a38..a5f22dbc89 100644 --- a/src/integrations/misc/jupyter-notebook-handler.ts +++ b/src/integrations/misc/jupyter-notebook-handler.ts @@ -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 { + static async fromFile(filePath: string, securityConfig?: JupyterSecurityConfig): Promise { 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 { + // 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) + } } diff --git a/src/integrations/misc/jupyter-notebook-security.ts b/src/integrations/misc/jupyter-notebook-security.ts new file mode 100644 index 0000000000..d48f0909bf --- /dev/null +++ b/src/integrations/misc/jupyter-notebook-security.ts @@ -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: /]*>[\s\S]*?<\/script>/gi, severity: "high" as const, type: "script_tag" }, + { pattern: /]*>/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: /]*>[\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): 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() + 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(/]*>[\s\S]*?<\/script>/gi, "[REMOVED: script tag]") + // Remove iframes + source = source.replace(/]*>[\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 (/): void { + this.config = { ...this.config, ...config } + } +}