10 |
11 |
12 |
13 | {{codebase_root}}
14 |
15 |
16 | {{codebase_basename}}
17 |
18 |
19 |
20 |
22 |
23 |
24 |
25 | {% include 'summary-card.html' %}
26 |
27 | {% include "errors.html" %}
28 |
29 | {% include 'issues.html' %}
30 |
31 | {% include "files-scanned-dialog.html" %}
32 |
33 | {% include 'footer.html' %}
34 |
35 | {% include "scripts.html" %}
36 |
--------------------------------------------------------------------------------
/nbdefense/templates/scripts.html:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/nbdefense/templates/summary-card.html:
--------------------------------------------------------------------------------
1 |
3 |
4 |
Files
5 | {{(artifacts_scanned["requirements"]|length) +
6 | (artifacts_scanned["notebooks"]|length)}}
7 |
8 |
9 |
Checks ran
10 | {{total_checks_run}}
11 |
12 |
13 |
Scanned at
14 | {{date}} · {{time}}
15 |
16 |
--------------------------------------------------------------------------------
/nbdefense/templating.py:
--------------------------------------------------------------------------------
1 | """
2 | This module loads Jinja2 templating engine.
3 | """
4 | from pathlib import Path
5 | from typing import Any, Optional, Union
6 |
7 | from jinja2 import Environment, FileSystemLoader, PackageLoader, Template
8 |
9 |
10 | class Jinja2Templates:
11 | def __init__(
12 | self, directory: Optional[Union[str, Path]] = None, **env_options: Any
13 | ) -> None:
14 | """
15 | Initialize templating environment.
16 | """
17 | self.env = self._create_env(directory, **env_options)
18 |
19 | def _create_env(
20 | self, directory: Optional[Union[str, Path]] = None, **env_options: Any
21 | ) -> Environment:
22 | loader: Union[PackageLoader, FileSystemLoader] = PackageLoader(
23 | "nbdefense", "templates"
24 | )
25 | if directory:
26 | directory = Path(directory)
27 | if directory.exists():
28 | loader = FileSystemLoader(directory)
29 |
30 | env_options.setdefault("loader", loader)
31 | env_options.setdefault("autoescape", True)
32 |
33 | env = Environment(**env_options) # nosec
34 | return env
35 |
36 | def get_template(self, name: str) -> Template:
37 | """
38 | Returns a template loaded via templating directory
39 |
40 | :param name: Name of the template e.g. report.html
41 | """
42 | return self.env.get_template(name)
43 |
--------------------------------------------------------------------------------
/nbdefense/tools.py:
--------------------------------------------------------------------------------
1 | """
2 | This module implements the tools to scan notebooks and repositories
3 | """
4 | import abc
5 | import os
6 | from subprocess import PIPE, Popen, run # nosec
7 | from typing import Any, List, Optional, Tuple, Union
8 |
9 | import click
10 |
11 |
12 | class Tool(metaclass=abc.ABCMeta):
13 | """
14 | A base class for different type of tools used in the scanning.
15 | """
16 |
17 | def __init__(self) -> None:
18 | pass
19 |
20 | @abc.abstractmethod
21 | def execute(
22 | self, args: Optional[Union[str, List[Any]]] = None
23 | ) -> Tuple[Union[int, Any], Optional[List[bytes]], Optional[List[bytes]]]:
24 | raise NotImplementedError
25 |
26 | @abc.abstractmethod
27 | def __call__(
28 | self,
29 | args: Optional[Union[str, List[Any]]] = None,
30 | ignore_output: Optional[bool] = True,
31 | ) -> Any:
32 | raise NotImplementedError
33 |
34 |
35 | class CLITool(Tool):
36 | """
37 | Executes shell commands for running scans using 3rd party tools.
38 | """
39 |
40 | def __init__(self) -> None:
41 | self.basecmd = ["echo"]
42 |
43 | def execute(
44 | self, args: Optional[Union[str, List[Any]]] = None
45 | ) -> Tuple[Union[int, Any], Optional[List[bytes]], Optional[List[bytes]]]:
46 | cmd = self.basecmd
47 | if args:
48 | if isinstance(args, List):
49 | cmd = self.basecmd + args
50 | else:
51 | cmd = self.basecmd + [args]
52 |
53 | p = Popen(cmd, stdout=PIPE, stderr=PIPE) # nosec
54 |
55 | stdout_lines = p.stdout.readlines() if p.stdout else None
56 | stderr_lines = p.stderr.readlines() if p.stderr else None
57 |
58 | return p.returncode, stdout_lines, stderr_lines
59 |
60 | def __call__(
61 | self,
62 | args: Optional[Union[str, List[Any]]] = None,
63 | ignore_output: Optional[bool] = True,
64 | ) -> Any:
65 | return self.execute(args)
66 |
67 |
68 | class Trivy(CLITool):
69 | def __init__(self, initialTrivyBinaryPath: str = "") -> None:
70 | super().__init__()
71 | self.installPath = os.path.split(__file__)[0]
72 | self.trivyPath = (
73 | initialTrivyBinaryPath
74 | if initialTrivyBinaryPath
75 | else os.path.join(self.installPath, "trivy")
76 | )
77 | self.basecmd = [self.trivyPath, "fs", "--security-checks", "vuln", "-f", "json"]
78 |
79 | def installed(self) -> bool:
80 | return os.path.isfile(self.trivyPath)
81 |
82 | def install(self) -> None:
83 | click.echo(f"Installing trivy to {self.trivyPath}")
84 | p1 = Popen(
85 | [
86 | "curl",
87 | "-sfL",
88 | "https://raw.githubusercontent.com/aquasecurity/trivy/main/contrib/install.sh",
89 | ],
90 | stdout=PIPE,
91 | ) # nosec
92 | p2 = run(
93 | ["sh", "-s", "--", "-b", self.installPath, "v0.32.1"],
94 | stdin=p1.stdout,
95 | stdout=PIPE,
96 | ) # nosec
97 |
98 |
99 | class DetectSecrets(CLITool):
100 | def __init__(self) -> None:
101 | super().__init__()
102 | self.basecmd = ["detect-secrets", "scan"]
103 |
--------------------------------------------------------------------------------
/nbdefense/utils.py:
--------------------------------------------------------------------------------
1 | import logging
2 | import os
3 | import re
4 | import socketserver
5 | import threading
6 | import time
7 | import webbrowser
8 | from datetime import datetime
9 | from http.server import SimpleHTTPRequestHandler
10 | from pathlib import Path
11 | from typing import Any, Optional
12 |
13 | import click
14 |
15 | logger = logging.getLogger(__name__)
16 |
17 |
18 | def serve_report_and_launch_url(report: str) -> None:
19 | """
20 | Create a HTTP server at a random port that servers the html passed as string
21 |
22 | :param report: The string that will be send back as response
23 |
24 | :return: TCPServer that can be launched with serve_forever()
25 | """
26 | encoded_report = report.encode("utf-8")
27 |
28 | class SimpleHTTPReportServer(SimpleHTTPRequestHandler):
29 | def do_GET(self) -> None:
30 | self.send_response(200)
31 | self.send_header("Content-type", "text/html")
32 | self.end_headers()
33 |
34 | self.wfile.write(encoded_report)
35 |
36 | handler = SimpleHTTPReportServer
37 | server = socketserver.TCPServer(("localhost", 0), handler)
38 | port = server.server_address[1]
39 | url = f"http://127.0.0.1:{port}/"
40 | click.echo(f"Report can be viewed at {url}")
41 |
42 | def start_server(server: Any) -> None:
43 | with server:
44 | server.serve_forever()
45 |
46 | server_thread = threading.Thread(target=start_server, args=(server,))
47 | server_thread.setDaemon(True)
48 | server_thread.start()
49 |
50 | webbrowser.open_new(url)
51 |
52 | should_run = True
53 |
54 | while should_run:
55 | try:
56 | time.sleep(1)
57 | except KeyboardInterrupt:
58 | should_run = False
59 |
60 |
61 | def write_output_file(report: str, output_format: str, output_file: str) -> Path:
62 | if output_file:
63 | output_file_path = Path(output_file)
64 | if output_file_path.parent and not output_file_path.parent.exists():
65 | os.makedirs(output_file_path.parent, exist_ok=True)
66 | else:
67 | # Pick file extension
68 | output_extension = None
69 | if output_format == "html":
70 | output_extension = ".html"
71 | elif output_format == "json":
72 | output_extension = ".json"
73 | else:
74 | raise NotImplementedError
75 |
76 | # Create file path
77 | now = datetime.now()
78 | date_time_str = now.strftime("%m%d-%H%M")
79 | file_name = "nbdefense" + date_time_str + output_extension
80 | output_file_path = Path(os.getcwd()) / file_name
81 |
82 | with open(output_file_path, "w", encoding="utf-8") as f:
83 | click.echo(report, file=f)
84 |
85 | return output_file_path
86 |
87 |
88 | def scrub_html(code_snippet: Optional[str]) -> str:
89 | if code_snippet:
90 | allowed_html_tags = ["table", "thead", "tbody", "tr", "th", "td"]
91 |
92 | new_code_snippet = code_snippet
93 | new_code_snippet = re.sub("<", "<", new_code_snippet)
94 | new_code_snippet = re.sub(">", ">", new_code_snippet)
95 |
96 | for allowed_html_tag in allowed_html_tags:
97 | # Tags without style information converted back to html tags
98 | new_code_snippet = re.sub(
99 | f"<{allowed_html_tag}>", f"<{allowed_html_tag}>", new_code_snippet
100 | )
101 | # Tags with style information converted back to html tags
102 | new_code_snippet = re.sub(
103 | f"<{allowed_html_tag}(.*)>",
104 | f"<{allowed_html_tag}\g<1>>",
105 | new_code_snippet,
106 | )
107 | new_code_snippet = re.sub(
108 | f"</{allowed_html_tag}>",
109 | f"",
110 | new_code_snippet,
111 | )
112 |
113 | return new_code_snippet
114 | else:
115 | return ""
116 |
--------------------------------------------------------------------------------
/pyproject.toml:
--------------------------------------------------------------------------------
1 | [tool.poetry]
2 | name = "nbdefense"
3 | description = "NB Defense CLI and SDK"
4 | version = "0.0.0"
5 | authors = ["ProtectAI
"]
6 | readme = "README.md"
7 | packages = [{ include = "nbdefense" }]
8 | license = "Apache License 2.0"
9 | exclude = ["tests/*", "Makefile"]
10 |
11 | [tool.poetry.scripts]
12 | nbdefense = "nbdefense.cli:cli"
13 |
14 | [tool.poetry.dependencies]
15 | python = ">=3.8,<3.11"
16 | click = "^8.1.3"
17 | requirements-parser = "0.5.0"
18 | detect-secrets = "1.4.0"
19 | nbformat = "5.6.1"
20 | presidio_analyzer = "2.2.29"
21 | nbconvert = "7.0.0"
22 | jinja2 = "^3.1.2"
23 | requests = "2.32.2"
24 | tqdm = "4.66.3"
25 | rich = "12.6.0"
26 | pandas = "^1.3.5"
27 | tomlkit = "0.11.6"
28 |
29 | [tool.poetry.group.test.dependencies]
30 | pytest = "7.2.0"
31 | pytest-watch = "4.2.0"
32 | responses = "0.22.0"
33 | bandit = { version = "1.7.5", extras = ["toml"] }
34 | mypy = "1.1.1"
35 | types-requests = "2.28.11.15"
36 | types-tqdm = "4.65.0.0"
37 | types-setuptools = "67.6.0.5"
38 |
39 | [tool.poetry.group.dev.dependencies]
40 | black = "24.3.0"
41 | pre-commit = "2.20.0"
42 | dunamai = "1.13.2"
43 | en_core_web_trf = [
44 | { url = "https://github.com/explosion/spacy-models/releases/download/en_core_web_trf-3.4.1/en_core_web_trf-3.4.1-py3-none-any.whl" },
45 | ]
46 |
47 | [tool.poetry.group.prod.dependencies]
48 | dunamai = "1.13.2"
49 |
50 | [tool.poetry-dynamic-versioning]
51 | enable = true
52 |
53 | [build-system]
54 | requires = ["poetry-core", "poetry-dynamic-versioning"]
55 | build-backend = "poetry.core.masonry.api"
56 |
57 | [tool.bandit]
58 | exclude_dirs = ["tests"]
59 |
--------------------------------------------------------------------------------
/tests/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/protectai/nbdefense/f836d3702636b715058e6e80f53165414ee159ff/tests/__init__.py
--------------------------------------------------------------------------------
/tests/conftest.py:
--------------------------------------------------------------------------------
1 | from tests.mock_notebooks.fixtures import (
2 | dataframe,
3 | dataframe_with_secrets,
4 | dataframe_with_custom_index,
5 | raw_notebook_cells,
6 | mock_notebook_as_json,
7 | mock_notebook,
8 | )
9 |
10 | from tests.plugin_tests.licenses.fixtures import (
11 | mock_license_cache,
12 | mock_fetch_licenses_from_dist_info_path,
13 | mock_fetch_license_data_from_pypi,
14 | create_license_requirements_file_path,
15 | )
16 |
17 | from tests.plugin_tests.cve.fixtures import (
18 | mock_third_party_dependencies,
19 | install_trivy,
20 | create_cve_requirements_file_path,
21 | )
22 |
--------------------------------------------------------------------------------
/tests/default_settings.py:
--------------------------------------------------------------------------------
1 | DEFAULT_SETTINGS = {
2 | "redact_secret": "PARTIAL",
3 | "trivy_binary_path": "",
4 | "plugins": {
5 | "nbdefense.plugins.SecretsPlugin": {
6 | "enabled": True,
7 | "secrets_plugins": [
8 | {"name": "SoftlayerDetector"},
9 | {"name": "StripeDetector"},
10 | {"name": "SendGridDetector"},
11 | {"name": "NpmDetector"},
12 | {"name": "KeywordDetector", "keyword_exclude": ""},
13 | {"name": "IbmCosHmacDetector"},
14 | {"name": "DiscordBotTokenDetector"},
15 | {"name": "BasicAuthDetector"},
16 | {"name": "AzureStorageKeyDetector"},
17 | {"name": "ArtifactoryDetector"},
18 | {"name": "AWSKeyDetector"},
19 | {"name": "CloudantDetector"},
20 | {"name": "GitHubTokenDetector"},
21 | {"name": "IbmCloudIamDetector"},
22 | {"name": "JwtTokenDetector"},
23 | {"name": "MailchimpDetector"},
24 | {"name": "PrivateKeyDetector"},
25 | {"name": "SlackDetector"},
26 | {"name": "SquareOAuthDetector"},
27 | {"name": "TwilioKeyDetector"},
28 | {"name": "Base64HighEntropyString", "limit": 4.5},
29 | {"name": "HexHighEntropyString", "limit": 3.0},
30 | ],
31 | },
32 | "nbdefense.plugins.PIIPlugin": {
33 | "enabled": True,
34 | "confidence_threshold": 0.8,
35 | "entities": {
36 | "US_PASSPORT": True,
37 | "AU_MEDICARE": True,
38 | "AU_TFN": True,
39 | "AU_ACN": True,
40 | "AU_ABN": True,
41 | "UK_NHS": True,
42 | "US_SSN": True,
43 | "US_ITIN": True,
44 | "US_DRIVER_LICENSE": True,
45 | "US_BANK_NUMBER": True,
46 | "MEDICAL_LICENSE": True,
47 | "LOCATION": True,
48 | "PHONE_NUMBER": True,
49 | "NRP": True,
50 | "IP_ADDRESS": True,
51 | "EMAIL_ADDRESS": True,
52 | "IBAN_CODE": True,
53 | "CRYPTO": True,
54 | "CREDIT_CARD": True,
55 | "PERSON": True,
56 | },
57 | },
58 | "nbdefense.plugins.LicenseDependencyFilePlugin": {
59 | "enabled": True,
60 | "accepted_licenses": [
61 | "Apache License 2.0",
62 | "Apache Software License",
63 | "Apache 2.0",
64 | "Apache-2.0",
65 | "BSD",
66 | "BSD License",
67 | "BSD 3-Clause",
68 | "BSD-3-Clause",
69 | "GNU Library or Lesser General Public License (LGPL)",
70 | "Microsoft Public License",
71 | "MIT",
72 | "MIT License",
73 | "Python Software Foundation License",
74 | "ISC License (ISCL)",
75 | "MIT-0",
76 | ],
77 | },
78 | "nbdefense.plugins.CVEDependencyFilePlugin": {"enabled": True},
79 | "nbdefense.plugins.LicenseNotebookPlugin": {
80 | "enabled": True,
81 | "accepted_licenses": [
82 | "Apache License 2.0",
83 | "Apache Software License",
84 | "Apache 2.0",
85 | "Apache-2.0",
86 | "BSD",
87 | "BSD License",
88 | "BSD 3-Clause",
89 | "BSD-3-Clause",
90 | "GNU Library or Lesser General Public License (LGPL)",
91 | "Microsoft Public License",
92 | "MIT",
93 | "MIT License",
94 | "Python Software Foundation License",
95 | "ISC License (ISCL)",
96 | "MIT-0",
97 | ],
98 | "licenses_for_notebooks_source": "HYBRID",
99 | },
100 | "nbdefense.plugins.CVENotebookPlugin": {"enabled": True},
101 | },
102 | }
103 |
--------------------------------------------------------------------------------
/tests/mock_notebooks/mock_notebook.py:
--------------------------------------------------------------------------------
1 | """
2 | Creates a mock v4 notebook using nbformat
3 | """
4 | from enum import Enum
5 | from typing import List, Optional
6 |
7 | import pandas as pd
8 | from nbformat import NotebookNode
9 | from nbformat.v4 import new_code_cell, new_markdown_cell, new_notebook, new_output
10 |
11 |
12 | class MockCellType(Enum):
13 | SOURCE = 1
14 | MARKDOWN = 2
15 |
16 |
17 | class MockNotebook:
18 | def __init__(self) -> None:
19 | self.nb = new_notebook(cells=[]) # type: ignore[no-untyped-call]
20 | self.current_cell_index = -1
21 |
22 | def add_cell(
23 | self,
24 | cell_type: MockCellType,
25 | source: List[str],
26 | stream_output: Optional[List[str]] = None,
27 | data_frame: Optional[pd.DataFrame] = None,
28 | html_output: Optional[List[str]] = None,
29 | plaintext_output: Optional[List[str]] = None,
30 | ) -> None:
31 | """
32 | Add a cell to the mock notebook.
33 | The outputs are add to source cell only. if cell type is markdown all other outputs are ignored.
34 | The function will add all the outputs to the cell in the order of function arguments.
35 | DataFrames are added as both html and plain-text format.
36 |
37 | :param cell_type: Type of the cell to add.
38 | :param source: Source or markdown string(s)
39 | :param stream_output:
40 | :param data_frame: Dataframe is converted to html using df.to_html()
41 | :param html_output:
42 | :param plaintext_output:
43 | """
44 | joined_source = "\n".join(source)
45 |
46 | joined_stream_output: Optional[str] = None
47 | if stream_output:
48 | joined_stream_output = "\n".join(stream_output)
49 |
50 | joined_html_output: Optional[str] = None
51 | if html_output:
52 | joined_html_output = "\n".join(html_output)
53 |
54 | joined_plaintext_output: Optional[str] = None
55 | if plaintext_output:
56 | joined_plaintext_output = "\n".join(plaintext_output)
57 |
58 | if cell_type == MockCellType.MARKDOWN:
59 | self.nb.cells.append(new_markdown_cell(joined_source)) # type: ignore[no-untyped-call]
60 | return
61 |
62 | self.nb.cells.append(new_code_cell(source=joined_source)) # type: ignore[no-untyped-call]
63 | if joined_stream_output:
64 | self.nb.cells[-1].outputs.append(
65 | new_output(output_type="stream", text=joined_stream_output) # type: ignore[no-untyped-call]
66 | )
67 |
68 | if data_frame is not None:
69 | self.nb.cells[-1].outputs.append(
70 | new_output(
71 | output_type="execute_result",
72 | data={
73 | "text/html": data_frame.to_html(),
74 | "text/plain": str(data_frame),
75 | },
76 | ) # type: ignore[no-untyped-call]
77 | )
78 |
79 | if joined_html_output:
80 | self.nb.cells[-1].outputs.append(
81 | new_output(
82 | output_type="execute_result", data={"text/html": joined_html_output}
83 | ) # type: ignore[no-untyped-call]
84 | )
85 |
86 | if joined_plaintext_output:
87 | self.nb.cells[-1].outputs.append(
88 | new_output(
89 | output_type="execute_result",
90 | data={"text/plain": joined_plaintext_output},
91 | ) # type: ignore[no-untyped-call]
92 | )
93 |
94 | @property
95 | def notebook_node(self) -> NotebookNode:
96 | return self.nb # type: ignore[no-any-return]
97 |
--------------------------------------------------------------------------------
/tests/plugin_tests/cve/fixtures.py:
--------------------------------------------------------------------------------
1 | import os
2 | from pathlib import Path
3 | from typing import Any
4 |
5 | import pytest
6 |
7 | from nbdefense.dependencies import DependencyInfo, ThirdPartyDependencies
8 | from nbdefense.plugins.cve.cve_plugin import CVEPlugin
9 |
10 |
11 | @pytest.fixture(scope="session")
12 | def create_cve_requirements_file_path() -> Path:
13 | current_directory = Path(os.path.dirname(os.path.realpath(__file__)))
14 | requirements_directory = current_directory / "mock_files"
15 | requirements_file_path = requirements_directory / "requirements.txt"
16 | with open(requirements_file_path, "w") as f:
17 | f.writelines(
18 | ["pytorch-lightning==1.5.10\n", "pandas==1.5.0\n", "matplotlibXtns==20.5\n"]
19 | )
20 | return requirements_file_path
21 |
22 |
23 | @pytest.fixture
24 | def mock_third_party_dependencies(monkeypatch: Any) -> None:
25 | dependencies = [
26 | DependencyInfo(
27 | name="pytorch-lightning",
28 | version="1.5.10",
29 | dist_info_path=Path("none"),
30 | dist_info_name="none",
31 | ),
32 | DependencyInfo(
33 | name="pandas",
34 | version="1.5.0",
35 | dist_info_path=Path("none"),
36 | dist_info_name="none",
37 | ),
38 | DependencyInfo(
39 | name="matplotlibXtns",
40 | version="20.5",
41 | dist_info_path=Path("none"),
42 | dist_info_name="none",
43 | ),
44 | ]
45 |
46 | def mock_parse_dependencies_available_in_env(self: Any) -> None:
47 | self._dependencies_in_env = dependencies
48 |
49 | def mock_parse_modules_from_dependencies_available(self: Any) -> None:
50 | self._dependencies_by_module["pytorch_lightning"] = dependencies[0]
51 | self._dependencies_by_module["pandas"] = dependencies[1]
52 | self._dependencies_by_module["matplotlibXtns"] = dependencies[2]
53 |
54 | monkeypatch.setattr(
55 | ThirdPartyDependencies,
56 | "_parse_dependencies_available_in_env",
57 | mock_parse_dependencies_available_in_env,
58 | )
59 | monkeypatch.setattr(
60 | ThirdPartyDependencies,
61 | "_parse_modules_from_dependencies_available",
62 | mock_parse_modules_from_dependencies_available,
63 | )
64 |
65 |
66 | @pytest.fixture
67 | def install_trivy() -> None:
68 | CVEPlugin.handle_binary_dependencies(True, True)
69 |
--------------------------------------------------------------------------------
/tests/plugin_tests/cve/mock_files/test-cve.ipynb:
--------------------------------------------------------------------------------
1 | {
2 | "cells": [
3 | {
4 | "cell_type": "code",
5 | "execution_count": 4,
6 | "id": "0d7a3ee3-2905-4755-b369-44144e9fa233",
7 | "metadata": {},
8 | "outputs": [
9 | {
10 | "name": "stdout",
11 | "output_type": "stream",
12 | "text": [
13 | "output\n"
14 | ]
15 | }
16 | ],
17 | "source": [
18 | "import pandas\n",
19 | "import matplotlibXtns\n",
20 | "client_secret = \"e2a6dda0953954cc1bf299b71e2d1befdd22709a\"\n",
21 | "print(\"output\")"
22 | ]
23 | },
24 | {
25 | "cell_type": "code",
26 | "execution_count": 3,
27 | "id": "26e1650c",
28 | "metadata": {},
29 | "outputs": [
30 | {
31 | "data": {
32 | "text/plain": [
33 | "'e2a6dda0953954cc1bf299b71e2d1befdd22709a'"
34 | ]
35 | },
36 | "execution_count": 3,
37 | "metadata": {},
38 | "output_type": "execute_result"
39 | }
40 | ],
41 | "source": [
42 | "import pytorch_lightning as lightning\n",
43 | "client_secret"
44 | ]
45 | }
46 | ],
47 | "metadata": {
48 | "kernelspec": {
49 | "display_name": "Python 3.8.13 ('nbdefense')",
50 | "language": "python",
51 | "name": "python3"
52 | },
53 | "language_info": {
54 | "codemirror_mode": {
55 | "name": "ipython",
56 | "version": 3
57 | },
58 | "file_extension": ".py",
59 | "mimetype": "text/x-python",
60 | "name": "python",
61 | "nbconvert_exporter": "python",
62 | "pygments_lexer": "ipython3",
63 | "version": "3.8.13"
64 | },
65 | "vscode": {
66 | "interpreter": {
67 | "hash": "7e2becf3a3002680eb4d2fb30f5037c726550c6b7f1911bce4ec05c40b2fbaa9"
68 | }
69 | }
70 | },
71 | "nbformat": 4,
72 | "nbformat_minor": 5
73 | }
74 |
--------------------------------------------------------------------------------
/tests/plugin_tests/cve/test_cve_dependency_file_plugin.py:
--------------------------------------------------------------------------------
1 | from pathlib import Path
2 | from typing import Any, List
3 |
4 | import pytest
5 |
6 | from nbdefense.codebase import Codebase
7 | from nbdefense.plugins.cve.cve_dependency_file_plugin import CVEDependencyFilePlugin
8 | from nbdefense.constants import DEFAULT_SETTINGS
9 |
10 |
11 | class TestCVEDependencyFilePlugin:
12 | @pytest.fixture
13 | def codebase(
14 | self, tmp_path: Path, create_cve_requirements_file_path: Path
15 | ) -> Codebase:
16 | return Codebase(
17 | Path("./tests/plugin_tests/cve/mock_files"),
18 | False,
19 | True,
20 | False,
21 | tmp_path,
22 | requirements_file=create_cve_requirements_file_path,
23 | )
24 |
25 | def test_scan(
26 | self,
27 | install_trivy: None,
28 | codebase: Codebase,
29 | create_cve_requirements_file_path: Path,
30 | ) -> None:
31 | CVEDependencyFilePlugin.scan(
32 | codebase,
33 | CVEDependencyFilePlugin.get_settings(
34 | "nbdefense.plugins.CVEDependencyFilePlugin", True, DEFAULT_SETTINGS
35 | ),
36 | )
37 | self.verify_scan_results(codebase, create_cve_requirements_file_path)
38 |
39 | @staticmethod
40 | def verify_scan_results(codebase: Codebase, requiremetns_file_path: Path) -> None:
41 | issues = codebase.issues.to_json()
42 | assert len(issues) == 2
43 |
44 | cve = "CVE-2022-0845"
45 | cve_filter = lambda issue: issue["details"]["results"]["VulnerabilityID"] == cve
46 | cve_list: List[Any] = list(filter(cve_filter, issues))
47 | assert len(cve_list) == 1
48 | cve1: Any = cve_list[0]
49 | assert cve1["code"] == "VULNERABLE_DEPENDENCY_DEP_FILE"
50 | assert cve1["severity"] == "CRITICAL"
51 | assert cve1["details"]["file_path"] == str(requiremetns_file_path)
52 | assert cve1["details"]["results"]["PkgName"] == "pytorch-lightning"
53 | assert cve1["details"]["results"]["InstalledVersion"] == "1.5.10"
54 | assert cve1["details"]["results"]["FixedVersion"] == "1.6.0"
55 |
56 | cve = "CVE-2021-4118"
57 | cve2: Any = list(filter(cve_filter, issues))
58 | assert len(cve2) == 1
59 | cve2 = cve2[0]
60 | assert cve2["code"] == "VULNERABLE_DEPENDENCY_DEP_FILE"
61 | assert cve2["severity"] == "HIGH"
62 | assert cve2["details"]["file_path"] == str(requiremetns_file_path)
63 | assert cve2["details"]["results"]["PkgName"] == "pytorch-lightning"
64 | assert cve2["details"]["results"]["InstalledVersion"] == "1.5.10"
65 | assert cve2["details"]["results"]["FixedVersion"] == "1.6.0"
66 |
--------------------------------------------------------------------------------
/tests/plugin_tests/cve/test_cve_notebook_plugin.py:
--------------------------------------------------------------------------------
1 | from pathlib import Path
2 | from typing import Any, List
3 |
4 | import pytest
5 |
6 | from nbdefense.codebase import Codebase
7 | from nbdefense.plugins.cve.cve_notebooks_plugin import CVENotebookPlugin
8 | from nbdefense.constants import DEFAULT_SETTINGS
9 |
10 |
11 | class TestCVENotebookPlugin:
12 | @pytest.fixture
13 | def codebase(self, mock_third_party_dependencies: None, tmp_path: Path) -> Codebase:
14 | return Codebase(
15 | Path("./tests/plugin_tests/cve/mock_files/test-cve.ipynb"),
16 | False,
17 | True,
18 | False,
19 | tmp_path,
20 | site_packages_path=Path("none"),
21 | )
22 |
23 | def test_scan(
24 | self,
25 | install_trivy: None,
26 | codebase: Codebase,
27 | ) -> None:
28 | CVENotebookPlugin.scan(
29 | codebase,
30 | CVENotebookPlugin.get_settings(
31 | "nbdefense.plugins.CVENotebookPlugin", True, DEFAULT_SETTINGS
32 | ),
33 | )
34 | self.verify_scan_results(codebase)
35 |
36 | @staticmethod
37 | def verify_scan_results(codebase: Codebase) -> None:
38 | notebooks = list(codebase.notebooks())
39 | assert len(notebooks) == 1
40 | issues = notebooks[0].issues.to_json()
41 | assert len(issues) == 2
42 |
43 | cve = "CVE-2022-0845"
44 | cve_filter = lambda issue: issue["details"]["summary_field"]["CVE_ID"] == cve
45 | cve_list: List[Any] = list(filter(cve_filter, issues))
46 | assert len(cve_list) == 1
47 | cve1: Any = cve_list[0]
48 | assert cve1["code"] == "VULNERABLE_DEPENDENCY_IMPORT"
49 | assert cve1["severity"] == "CRITICAL"
50 | assert cve1["cell"]["cell_index"] == 1
51 | assert cve1["line_index"] == 0
52 | assert cve1["location"] == "INPUT"
53 | assert cve1["character_start_index"] == 7
54 | assert cve1["character_end_index"] == 24
55 | assert (
56 | cve1["details"]["summary_field"]["INSTALLED_PACKAGE"] == "pytorch-lightning"
57 | )
58 | assert cve1["details"]["summary_field"]["INSTALLED_VERSION"] == "1.5.10"
59 | assert cve1["details"]["summary_field"]["FIXED_VERSION"] == "1.6.0"
60 |
61 | cve = "CVE-2021-4118"
62 | cve2: Any = list(filter(cve_filter, issues))
63 | assert len(cve2) == 1
64 | cve2 = cve2[0]
65 | assert cve2["code"] == "VULNERABLE_DEPENDENCY_IMPORT"
66 | assert cve2["severity"] == "HIGH"
67 | assert cve1["cell"]["cell_index"] == 1
68 | assert cve1["line_index"] == 0
69 | assert cve1["location"] == "INPUT"
70 | assert cve1["character_start_index"] == 7
71 | assert cve1["character_end_index"] == 24
72 | assert (
73 | cve2["details"]["summary_field"]["INSTALLED_PACKAGE"] == "pytorch-lightning"
74 | )
75 | assert cve2["details"]["summary_field"]["INSTALLED_VERSION"] == "1.5.10"
76 | assert cve2["details"]["summary_field"]["FIXED_VERSION"] == "1.6.0"
77 |
--------------------------------------------------------------------------------
/tests/plugin_tests/licenses/fixtures.py:
--------------------------------------------------------------------------------
1 | import os
2 | from pathlib import Path
3 | from typing import Any, Dict, List
4 |
5 | import pytest
6 | import requests
7 |
8 | from nbdefense.codebase import PackageInfo
9 | from nbdefense.plugins.licenses.license_plugin import LicenseCache, LicensePlugin
10 |
11 |
12 | @pytest.fixture(scope="session")
13 | def create_license_requirements_file_path() -> Path:
14 | current_directory = Path(os.path.dirname(os.path.realpath(__file__)))
15 | requirements_directory = current_directory / "mock_files"
16 | requirements_file_path = requirements_directory / "requirements.txt"
17 | with open(requirements_file_path, "w") as f:
18 | f.writelines(
19 | [
20 | "test_package_1==23.2.1\n",
21 | "test_package_2==2.28.1\n",
22 | "test_package_3==1.8.0\n",
23 | "test_package_4==1.16.0",
24 | ]
25 | )
26 | return requirements_file_path
27 |
28 |
29 | @pytest.fixture
30 | def mock_license_cache(monkeypatch: Any) -> None:
31 | def mock_init(self: Any, temp_directory: Path, filename: str) -> None:
32 | self._cache = {}
33 | self._cache_updated = False
34 |
35 | def mock_del(self: Any) -> None:
36 | pass
37 |
38 | monkeypatch.setattr(LicenseCache, "__init__", mock_init)
39 | monkeypatch.setattr(LicenseCache, "__del__", mock_del)
40 |
41 |
42 | @pytest.fixture
43 | def mock_fetch_licenses_from_dist_info_path(monkeypatch: Any) -> Dict[str, List[str]]:
44 | classifiers: List[str] = []
45 | license_metadata: List[str] = []
46 | package_names: List[str] = []
47 |
48 | def mock_parse_license_from_dist_info_path(package: PackageInfo) -> Any:
49 | if not package_names or package.name in package_names:
50 | if classifiers:
51 | return classifiers
52 | else:
53 | return license_metadata
54 | return []
55 |
56 | monkeypatch.setattr(
57 | LicensePlugin,
58 | "_parse_license_data_from_dist_info_path",
59 | mock_parse_license_from_dist_info_path,
60 | )
61 | return {
62 | "classifiers": classifiers,
63 | "license_metadata": license_metadata,
64 | "package_names": package_names,
65 | }
66 |
67 |
68 | @pytest.fixture
69 | def mock_fetch_license_data_from_pypi(monkeypatch: Any) -> Dict[str, List[str]]:
70 | licenses: List[str] = []
71 | package_names: List[str] = []
72 |
73 | def mock_function_response(
74 | _: requests.Session, package: PackageInfo
75 | ) -> PackageInfo:
76 | if not package_names or package.name in package_names:
77 | package.licenses = licenses
78 | else:
79 | package.licenses = []
80 | return package
81 |
82 | monkeypatch.setattr(
83 | LicensePlugin,
84 | "_fetch_license_data_from_pypi",
85 | mock_function_response,
86 | )
87 | return {
88 | "licenses": licenses,
89 | "package_names": package_names,
90 | }
91 |
--------------------------------------------------------------------------------
/tests/plugin_tests/licenses/mock_files/test-license.ipynb:
--------------------------------------------------------------------------------
1 | {
2 | "cells": [
3 | {
4 | "cell_type": "code",
5 | "execution_count": 4,
6 | "id": "0d7a3ee3-2905-4755-b369-44144e9fa233",
7 | "metadata": {},
8 | "outputs": [
9 | {
10 | "name": "stdout",
11 | "output_type": "stream",
12 | "text": [
13 | "output\n"
14 | ]
15 | }
16 | ],
17 | "source": [
18 | "import pandas\n",
19 | "import matplotlibXtns\n",
20 | "client_secret = \"e2a6dda0953954cc1bf299b71e2d1befdd22709a\"\n",
21 | "print(\"output\")"
22 | ]
23 | },
24 | {
25 | "cell_type": "code",
26 | "execution_count": 3,
27 | "id": "26e1650c",
28 | "metadata": {},
29 | "outputs": [
30 | {
31 | "data": {
32 | "text/plain": [
33 | "'e2a6dda0953954cc1bf299b71e2d1befdd22709a'"
34 | ]
35 | },
36 | "execution_count": 3,
37 | "metadata": {},
38 | "output_type": "execute_result"
39 | }
40 | ],
41 | "source": [
42 | "import pytorch_lightning as lightning\n",
43 | "client_secret"
44 | ]
45 | }
46 | ],
47 | "metadata": {
48 | "kernelspec": {
49 | "display_name": "Python 3.8.13 ('nbdefense')",
50 | "language": "python",
51 | "name": "python3"
52 | },
53 | "language_info": {
54 | "codemirror_mode": {
55 | "name": "ipython",
56 | "version": 3
57 | },
58 | "file_extension": ".py",
59 | "mimetype": "text/x-python",
60 | "name": "python",
61 | "nbconvert_exporter": "python",
62 | "pygments_lexer": "ipython3",
63 | "version": "3.8.13"
64 | },
65 | "vscode": {
66 | "interpreter": {
67 | "hash": "7e2becf3a3002680eb4d2fb30f5037c726550c6b7f1911bce4ec05c40b2fbaa9"
68 | }
69 | }
70 | },
71 | "nbformat": 4,
72 | "nbformat_minor": 5
73 | }
74 |
--------------------------------------------------------------------------------
/tests/plugin_tests/licenses/test_license_dependency_file_plugin.py:
--------------------------------------------------------------------------------
1 | from pathlib import Path
2 | from typing import Dict, List
3 |
4 | import pytest
5 |
6 | from nbdefense.codebase import Codebase
7 | from nbdefense.plugins.licenses.licenses_dependency_file_plugin import (
8 | LicenseDependencyFilePlugin,
9 | )
10 | from tests.default_settings import DEFAULT_SETTINGS
11 |
12 |
13 | class TestLicenseDependencyFilePlugin:
14 | @pytest.fixture
15 | def codebase(
16 | self, tmp_path: Path, create_license_requirements_file_path: Path
17 | ) -> Codebase:
18 | print(create_license_requirements_file_path)
19 | return Codebase(
20 | Path("./tests/plugin_tests/licenses/mock_files"),
21 | False,
22 | True,
23 | False,
24 | tmp_path,
25 | requirements_file=create_license_requirements_file_path,
26 | )
27 |
28 | def test_scan(
29 | self,
30 | mock_license_cache: None,
31 | mock_fetch_licenses_from_dist_info_path: Dict[str, List[str]],
32 | mock_fetch_license_data_from_pypi: Dict[str, List[str]],
33 | codebase: Codebase,
34 | create_license_requirements_file_path: Path,
35 | ) -> None:
36 | unapproved_licenses = ["BAD_LICENSE_1", "BAD_LICENSE_2"]
37 | mock_fetch_license_data_from_pypi["licenses"].append(unapproved_licenses[0])
38 | mock_fetch_license_data_from_pypi["licenses"].append(unapproved_licenses[1])
39 | LicenseDependencyFilePlugin.scan(
40 | codebase,
41 | LicenseDependencyFilePlugin.get_settings(
42 | plugin_class_name="nbdefense.plugins.LicenseDependencyFilePlugin",
43 | settings=DEFAULT_SETTINGS,
44 | ),
45 | )
46 | self.verify_scan_results(codebase, create_license_requirements_file_path)
47 |
48 | @staticmethod
49 | def verify_scan_results(codebase: Codebase, requirements_file_path: Path) -> None:
50 | issues = codebase.issues.to_json()
51 | assert len(issues) == 8
52 | for package_name_index in range(1, 5):
53 | package_name = f"test_package_{package_name_index}"
54 | package_version = None
55 | if package_name == "test_package_1":
56 | package_version = "23.2.1"
57 | elif package_name == "test_package_2":
58 | package_version = "2.28.1"
59 | elif package_name == "test_package_3":
60 | package_version = "1.8.0"
61 | elif package_name == "test_package_4":
62 | package_version = "1.16.0"
63 |
64 | issues_filter_by_package_name = list(
65 | filter(
66 | lambda issue: (issue["details"]["package_name"] == package_name),
67 | issues,
68 | )
69 | )
70 | assert len(issues_filter_by_package_name) == 2
71 | for package_index, _ in enumerate(issues_filter_by_package_name):
72 | assert {
73 | "code": "UNAPPROVED_LICENSE_DEP_FILE",
74 | "severity": "MEDIUM",
75 | "details": {
76 | "package_name": package_name,
77 | "package_version": package_version,
78 | "unapproved_license": f"BAD_LICENSE_{package_index+1}",
79 | "file_path": str(requirements_file_path),
80 | },
81 | } in issues_filter_by_package_name
82 |
--------------------------------------------------------------------------------
/tests/plugin_tests/licenses/test_license_plugin.py:
--------------------------------------------------------------------------------
1 | from pathlib import Path
2 | from typing import Dict, List
3 |
4 | from nbdefense.codebase import PackageInfo
5 | from nbdefense.plugins.licenses.license_plugin import LicensePlugin
6 | from nbdefense.plugins.licenses.license_plugin_settings import LicensePluginSource
7 | from tests.default_settings import DEFAULT_SETTINGS
8 |
9 |
10 | class TestLicensePlugin:
11 | def test_get_licenses_none_found(
12 | self,
13 | tmp_path: Path,
14 | mock_license_cache: None,
15 | mock_fetch_licenses_from_dist_info_path: Dict[str, List[str]],
16 | mock_fetch_license_data_from_pypi: Dict[str, List[str]],
17 | ) -> None:
18 | test_package = PackageInfo(name="none", version="none")
19 |
20 | installed_packages = LicensePlugin.get_licenses(
21 | [test_package], tmp_path, LicensePluginSource.HYBRID
22 | )
23 | assert len(installed_packages) == 1
24 | assert installed_packages[0].name == test_package.name
25 | assert installed_packages[0].version == test_package.version
26 | assert not installed_packages[0].licenses
27 |
28 | def test_get_licenses_fall_back_to_api(
29 | self,
30 | tmp_path: Path,
31 | mock_license_cache: None,
32 | mock_fetch_licenses_from_dist_info_path: Dict[str, List[str]],
33 | mock_fetch_license_data_from_pypi: Dict[str, List[str]],
34 | ) -> None:
35 | package_info = PackageInfo(name="package_test_1", version="version")
36 | fallback_license = "Fallback To API License"
37 | mock_fetch_license_data_from_pypi["licenses"].append(fallback_license)
38 | installed_packages = LicensePlugin.get_licenses(
39 | [package_info], tmp_path, LicensePluginSource.HYBRID
40 | )
41 | assert len(installed_packages) == 1
42 | assert installed_packages[0].name == package_info.name
43 | assert installed_packages[0].version == package_info.version
44 | assert installed_packages[0].licenses == [fallback_license]
45 |
46 | def test_filter_for_unapproved_licenses(self) -> None:
47 | licenses = LicensePlugin.filter_for_unapproved_licenses(
48 | ["MIT", "UNAPPROVED"], ["MIT"]
49 | )
50 | assert len(licenses) == 1
51 | assert licenses[0] == "UNAPPROVED"
52 |
53 | def test_get_licenses_in_metadata(
54 | self,
55 | tmp_path: Path,
56 | mock_license_cache: None,
57 | mock_fetch_licenses_from_dist_info_path: Dict[str, List[str]],
58 | ) -> None:
59 | license = "GNU"
60 | package_info = PackageInfo(name="package_test_1", version="version")
61 | mock_fetch_licenses_from_dist_info_path["license_metadata"].append(license)
62 | installed_packages = LicensePlugin.get_licenses(
63 | [package_info], tmp_path, LicensePluginSource.HYBRID
64 | )
65 | assert len(installed_packages) == 1
66 | assert installed_packages[0].name == package_info.name
67 | assert installed_packages[0].version == package_info.version
68 | assert installed_packages[0].licenses == [license]
69 |
70 | def test_process_for_unnapproved_licenses(self) -> None:
71 | settings = LicensePlugin.get_settings(
72 | plugin_class_name="nbdefense.plugins.LicenseNotebookPlugin",
73 | settings=DEFAULT_SETTINGS,
74 | )
75 | unapproved_license = "APACHE_UNAPPROVED"
76 | unapproved_licenses = LicensePlugin.filter_for_unapproved_licenses(
77 | ["MIT", unapproved_license], settings.get_accepted_licenses()
78 | )
79 | assert unapproved_licenses == [unapproved_license]
80 |
81 | def test_parse_licenses_from_classifiers_single_license(self) -> None:
82 | classifiers_list = [
83 | "License :: OSI Approved",
84 | "License :: OSI Approved :: MIT",
85 | "Private :: Do Not Upload",
86 | ]
87 | licenses = LicensePlugin._parse_licenses_from_classifiers(classifiers_list)
88 | assert licenses == ["MIT"]
89 |
90 | def test_parse_licenses_from_classifiers_multiple_license(self) -> None:
91 | classifiers_list = [
92 | "License :: OSI Approved :: MIT",
93 | "License :: OSI Approved :: Apache 2",
94 | ]
95 | licenses = LicensePlugin._parse_licenses_from_classifiers(classifiers_list)
96 | assert licenses == ["MIT", "Apache 2"]
97 |
98 | def test_parse_licenses_from_classifiers_no_license(self) -> None:
99 | no_license_classifier = ["Private :: Do Not Upload"]
100 | licenses = LicensePlugin._parse_licenses_from_classifiers(no_license_classifier)
101 | assert licenses == []
102 |
--------------------------------------------------------------------------------
/tests/test_codebase.py:
--------------------------------------------------------------------------------
1 | import pathlib
2 |
3 | import pytest
4 |
5 | from nbdefense.codebase import Codebase
6 |
7 |
8 | class TestCodebase:
9 | @pytest.fixture
10 | def path_to_single_notebook(
11 | self, tmp_path: pathlib.Path, mock_notebook_as_json: str
12 | ) -> pathlib.Path:
13 | tmp_file = tmp_path.joinpath("temp.ipynb")
14 | with open(tmp_file, "w") as f:
15 | f.write(mock_notebook_as_json)
16 |
17 | return tmp_file
18 |
19 | @pytest.fixture
20 | def path_with_multiple_notebooks(
21 | self, tmp_path: pathlib.Path, mock_notebook_as_json: str
22 | ) -> pathlib.Path:
23 | paths = ["temp.ipynb", "dir2/temp2.ipynb", "dir3/temp3.ipynb"]
24 | for path in paths:
25 | tmp_file = tmp_path.joinpath(path)
26 | if "/" in path:
27 | subdir_path = tmp_path.joinpath(path.split("/")[0])
28 | subdir_path.mkdir(exist_ok=True)
29 | with open(tmp_file, "w") as f:
30 | f.write(mock_notebook_as_json)
31 |
32 | return tmp_path
33 |
34 | def test_codebase_single_notebook_file(
35 | self, path_to_single_notebook: pathlib.Path, tmp_path: pathlib.Path
36 | ) -> None:
37 | # TODO: Add requirements.txt path as well.
38 | c = Codebase(
39 | path=path_to_single_notebook,
40 | recursive=False,
41 | quiet=True,
42 | temp_directory=tmp_path,
43 | requirements_file=None,
44 | show_progress_bars=False,
45 | )
46 |
47 | assert c.quiet == True
48 | assert len(list(c.notebooks())) == 1
49 |
50 | def test_codebase_multiple_notebook_path(
51 | self, path_with_multiple_notebooks: pathlib.Path, tmp_path: pathlib.Path
52 | ) -> None:
53 | c = Codebase(
54 | path=path_with_multiple_notebooks,
55 | recursive=False,
56 | quiet=True,
57 | temp_directory=tmp_path,
58 | requirements_file=None,
59 | show_progress_bars=False,
60 | )
61 |
62 | assert len(list(c.notebooks())) == 1
63 |
64 | c = Codebase(
65 | path=path_with_multiple_notebooks,
66 | recursive=True,
67 | quiet=True,
68 | temp_directory=tmp_path,
69 | requirements_file=None,
70 | show_progress_bars=False,
71 | )
72 |
73 | assert len(list(c.notebooks())) == 3
74 |
--------------------------------------------------------------------------------
/tests/test_notebook.py:
--------------------------------------------------------------------------------
1 | import pathlib
2 | from typing import Dict, List, Optional
3 |
4 | import pandas as pd
5 |
6 | from nbdefense.notebook import (
7 | Cell,
8 | InputCell,
9 | InputCellType,
10 | Notebook,
11 | OutputCell,
12 | OutputCellType,
13 | )
14 |
15 |
16 | class TestNotebookParsing:
17 | def test_parsing(
18 | self, mock_notebook: Notebook, raw_notebook_cells: List[Dict[str, List[str]]]
19 | ) -> None:
20 | assert len(mock_notebook.input_cells) == len(raw_notebook_cells)
21 | assert (
22 | len(mock_notebook.output_cells) == 8
23 | ) # TODO replace the hard-coded output number
24 |
25 | for index, data in enumerate(raw_notebook_cells):
26 | cell = mock_notebook.input_cells[index]
27 | assert type(cell) == InputCell
28 | if "markdown" in data:
29 | assert cell.cell_type == InputCellType.MARKDOWN
30 | markdown_text = data["markdown"]
31 | assert "\n".join(markdown_text) == str(cell)
32 | else:
33 | assert cell.cell_type == InputCellType.SOURCE
34 | source_text = data["source"]
35 | assert "\n".join(source_text) == str(cell)
36 |
37 | index = 0
38 | for data in raw_notebook_cells:
39 | if "plaintext_output" in data:
40 | cell = mock_notebook.output_cells[index]
41 | assert type(cell) == OutputCell
42 | assert cell.cell_type == OutputCellType.PLAINTEXT
43 | assert "\n".join(data["plaintext_output"]) == str(cell)
44 | index = index + 1
45 | elif "stream_output" in data:
46 | cell = mock_notebook.output_cells[index]
47 | assert type(cell) == OutputCell
48 | assert cell.cell_type == OutputCellType.STREAM
49 | assert "\n".join(data["stream_output"]) == str(cell)
50 | index = index + 1
51 | elif "dataframe_output" in data:
52 | cell = mock_notebook.output_cells[index]
53 | assert cell.cell_type == OutputCellType.DATAFRAME
54 | df: pd.DataFrame = data["dataframe_output"]
55 | assert [
56 | str(elem) for elem in list(df.columns.values)
57 | ] == cell.data_cells[ # type: ignore[attr-defined]
58 | 0
59 | ]
60 | idx = 0
61 | for _, row in df.iterrows():
62 | assert [str(val) for val in list(row.values)] == cell.data_cells[ # type: ignore[attr-defined]
63 | idx + 1
64 | ]
65 | idx = idx + 1
66 |
67 | index = index + 1
68 |
69 | def test_input_file_line_number_to_cell(
70 | self, tmp_path: pathlib.Path, mock_notebook: Notebook
71 | ) -> None:
72 | tmp_file = tmp_path.joinpath("input_test_file.py")
73 | with open(tmp_file, "w") as f:
74 | for cell in mock_notebook.input_cells:
75 | f.write(str(cell))
76 | f.write("\n")
77 |
78 | with open(tmp_file, "r") as f:
79 | for line_number, line in enumerate(f.readlines()):
80 | line_number = line_number + 1
81 | found_cell: Optional[
82 | Cell
83 | ] = mock_notebook.get_input_file_line_number_to_cell(line_number)
84 | assert found_cell
85 | # Test for the file text is in the cell found at line number
86 | assert line[:-1] in found_cell.lines
87 |
88 | # Test for the file text at line index within the found cell
89 | cell_line_index = found_cell.get_cell_line_index_for_file_line_number(
90 | line_number
91 | )
92 | assert cell_line_index is not None
93 | assert len(found_cell.lines) > cell_line_index
94 | assert found_cell.lines[cell_line_index] == line[:-1]
95 |
96 | def test_output_file_line_number_to_cell(
97 | self, tmp_path: pathlib.Path, mock_notebook: Notebook
98 | ) -> None:
99 | tmp_file = tmp_path.joinpath("output_test_file.txt")
100 | with open(tmp_file, "w") as f:
101 | for cell in mock_notebook.output_cells:
102 | f.write(str(cell))
103 | f.write("\n")
104 |
105 | with open(tmp_file, "r") as f:
106 | for line_number, line in enumerate(f.readlines()):
107 | line_number = line_number + 1
108 | found_cell = mock_notebook.get_output_file_line_number_to_cell(
109 | line_number
110 | )
111 |
112 | # Test for the file text is in the cell found at line number
113 | assert found_cell
114 | assert line[:-1] in found_cell.lines
115 |
116 | # Test for the file text at line index within the found cell
117 | cell_line_index = found_cell.get_cell_line_index_for_file_line_number(
118 | line_number
119 | )
120 | assert cell_line_index is not None
121 | assert found_cell.lines[cell_line_index] == line[:-1]
122 |
--------------------------------------------------------------------------------