diff --git a/README.md b/README.md index 9da86e5e..73539589 100644 --- a/README.md +++ b/README.md @@ -53,6 +53,7 @@ Metis includes support for the following languages: | TableGen | Tools | Built-in plugin | | Verilog | Tree-sitter + Structural Analysis + tools| Built-in plugin | | AArch64 Assembly | Tools | Built-in plugin | +| Jupyter NB | Tree-sitter + Structural Analysis + tools| Built-in plugin | For triage analysis details (`Flow Analysis` vs `Structural Analysis`), see [docs/triage-flow.md](docs/triage-flow.md). diff --git a/src/metis/plugins/ipynb_plugin.py b/src/metis/plugins/ipynb_plugin.py new file mode 100644 index 00000000..8db3c47f --- /dev/null +++ b/src/metis/plugins/ipynb_plugin.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright 2025 Arm Limited and/or its affiliates +# SPDX-License-Identifier: Apache-2.0 + +import json +from llama_index.core.node_parser import CodeSplitter +from llama_index.core.schema import Document + +from metis.plugins.base import BaseLanguagePlugin + + +class IpynbPlugin(BaseLanguagePlugin): + name = "ipynb" + + def __init__(self, plugin_config): + self.plugin_config = plugin_config + + def get_name(self): + return "ipynb" + + def can_handle(self, extension): + return extension.lower() == ".ipynb" + + def get_supported_extensions(self): + return [".ipynb"] + + def get_splitter(self): + splitting_cfg = ( + self.plugin_config.get("plugins", {}) + .get(self.get_name(), {}) + .get("splitting", {}) + ) + return NotebookCodeSplitter( + chunk_lines=splitting_cfg["chunk_lines"], + chunk_lines_overlap=splitting_cfg["chunk_lines_overlap"], + max_chars=splitting_cfg["max_chars"], + ) + + def get_prompts(self): + return ( + self.plugin_config.get("plugins", {}) + .get(self.get_name(), {}) + .get("prompts", {}) + ) + + +class NotebookCodeSplitter(CodeSplitter): + def __init__(self, **kwargs): + super().__init__(language="python", **kwargs) + + def get_nodes_from_documents(self, documents, show_progress=False, **kwargs): + processed_docs = [] + for doc in documents: + processed_docs.append(self._extract_notebook_code(doc)) + return super().get_nodes_from_documents(processed_docs, show_progress, **kwargs) + + def _extract_notebook_code(self, doc): + try: + notebook = json.loads(doc.text) + code_cells = [] + for i, cell in enumerate(notebook.get("cells", [])): + if cell.get("cell_type") == "code": + source = "".join(cell.get("source", [])) + if source.strip(): + code_cells.append(f"# Cell {i + 1}\n{source}") + extracted_code = "\n\n".join(code_cells) + print("......") + return Document(text=extracted_code, metadata=doc.metadata, id_=doc.id_) + except (json.JSONDecodeError, KeyError): + return doc diff --git a/src/metis/plugins/languages/ipynb.yaml b/src/metis/plugins/languages/ipynb.yaml new file mode 100644 index 00000000..8db06d80 --- /dev/null +++ b/src/metis/plugins/languages/ipynb.yaml @@ -0,0 +1,58 @@ +splitting: + chunk_lines: 40 + chunk_lines_overlap: 15 + max_chars: 1500 +prompts: + security_review: |- + You are a thorough security engineer specializing in Python and Jupyter notebooks. + Always tie every finding directly to evidence in FILE_CHANGES, RELEVANT_CONTEXT, and ORIGINAL_FILE. + Do not report security issues that are not supported by the provided code or context. + You will be given: + 1. FILE_CHANGES - Code changes where "+" indicates added lines and "-" indicates removed lines. + 2. RELEVANT_CONTEXT - Additional information describing the purpose of the changes. + 3. ORIGINAL_FILE - The original file before modification. This may be empty. + Your tasks are: + 1. Security Review Scope + - Review the security implications of FILE_CHANGES, focusing primarily on added ("+") and removed ("-") lines while considering how they interact with the rest of the notebook. + security_review_checks: |- + 2. What to Check + - Review for security issues including, but not limited to: + - OWASP Top 10 vulnerabilities + - Hardcoded secrets, credentials, tokens, API keys, URLs, or endpoints + - Insecure use of third-party libraries + - Unsafe deserialization or dynamic code execution (eval, exec, pickle, etc.) + - SQL, shell, or command injection risks + - Exposure of sensitive data in notebook outputs or saved cell results + - Insecure notebook execution patterns + - Unsafe use of Databricks utilities (dbutils, secrets, widgets, filesystem access) + - Missing or improper use of Databricks Secrets + - Imports that unnecessarily increase the attack surface + - Unsafe cell execution order assumptions + - Logging or printing sensitive information + - Weak authentication or authorization patterns + - Unsafe filesystem or cloud storage access + - Include the jupyter notebook metadata in the review to check for any security implications. + - Do not report code quality, style, performance, or maintainability issues unless they directly impact security. + validation_review: |- + Validate the following Jupyter notebook security review. + You will be given: + SNIPPET: The relevant Jupyter notebook code snippet. + REVIEW: A list of potential security issues identified in the code changes. + Your tasks are: + 1. Examine each item in REVIEW and verify whether it represents a genuine security issue based on the SNIPPET. + 2. Remove issues that are false positives or are already mitigated. + 3. Keep only issues that represent real security risks. + 4. Improve the remaining findings by adding missing technical details or security impact where necessary. + 5. If no valid security issues remain, respond with an empty array ([]). + snippet_security_summary: "Summarize the security implications of these Jupyter notebook code changes." + attempt_fix: "Based on the issues detected in the Jupyter notebook code changes, propose a fix patch. Issues: {issues} Patch: {patch}" + security_review_file: |- + You are a thorough security engineer specializing in Python and Jupyter notebooks. + Always tie every finding directly to evidence in FILE and CONTEXT. + Do not report security issues that are not supported by the provided code or context. + You will be given: + 1. FILE - A Jupyter notebook or Python source file. + 2. CONTEXT - Additional information describing the file. + Your tasks are: + 1. Review the entire file for security issues. + 2. Report only actionable security findings with supporting evidence from the code. diff --git a/src/metis/plugins/manifests/ipynb.yaml b/src/metis/plugins/manifests/ipynb.yaml new file mode 100644 index 00000000..7c53072d --- /dev/null +++ b/src/metis/plugins/manifests/ipynb.yaml @@ -0,0 +1,10 @@ +name: ipynb +aliases: +- ipynb +extensions: +- .ipynb +filename_patterns: [] +implementation: metis.plugins.ipynb_plugin:IpynbPlugin +config_resource: languages/ipynb.yaml +capabilities: {} +priority: 0 diff --git a/tests/test_language_plugin_registry.py b/tests/test_language_plugin_registry.py index 9df63629..51c19263 100644 --- a/tests/test_language_plugin_registry.py +++ b/tests/test_language_plugin_registry.py @@ -329,13 +329,13 @@ def test_registry_loads_required_prompt_keys_for_supported_languages(): ] ) } - assert registry.supported_language_names() == [ "aarch64_assembly", "c", "cpp", "csharp", "go", + "ipynb", "java", "javascript", "kotlin",