From d8231755da2bd1edd798e7e0062dc509b0a0922e Mon Sep 17 00:00:00 2001 From: Anil Balireddy Date: Fri, 11 Sep 2026 10:05:58 -0700 Subject: [PATCH] Harden DownloadSecFilingsStage: validate_config + injection-safe command building DownloadSecFilingsStage.execute() had two gaps relative to the rest of the finance recipe's stages: 1. No validate_config() override, so a workflow YAML missing output_dir/sec_identity_email/sec_identity_company/tickers/ start_year/end_year raised a bare KeyError with no indication of which field or file to fix. Add validate_config() following the pattern used elsewhere in this recipe (e.g. train_validation_split.py), including the "config" file alt-path where tickers/start_year/end_year come from a separate YAML instead. 2. The shell command was built via raw f-string interpolation with manual double-quotes instead of nvflow.lib.cli_cmd.build_python_cmd, the shlex.quote-based helper this repo already uses for every other stage that submits a command this way (validate_questions, data_transformation, apply_prompt_template, convert_to_responses_api, prepare_data, prefetch_cache). A ticker or company name containing backticks or $(...) would not have been safely escaped under the old f-string approach. Migrate to build_python_cmd, which also standardizes on python3 (this repo's documented interpreter convention) instead of the unversioned python entry point the old code used. Add tests/test_download_sec_filings.py: validate_config coverage for every required field (direct-config and config-file-reference paths), plus execute() tests that stub nemo_skills.pipeline.cli in sys.modules (the module doesn't import nemo_skills at module level, so this works in the lightweight CI environment) to capture the exact rendered command -- including a regression test asserting that shell metacharacters in a ticker value are safely single-quoted rather than interpreted by the shell. Co-Authored-By: Claude Sonnet 5 Signed-off-by: Anil Balireddy --- .../stages/download/download_sec_filings.py | 37 ++-- tests/test_download_sec_filings.py | 158 ++++++++++++++++++ 2 files changed, 185 insertions(+), 10 deletions(-) create mode 100644 tests/test_download_sec_filings.py diff --git a/nvflow/recipes/finance/stages/download/download_sec_filings.py b/nvflow/recipes/finance/stages/download/download_sec_filings.py index 8c1942f..2802fb3 100644 --- a/nvflow/recipes/finance/stages/download/download_sec_filings.py +++ b/nvflow/recipes/finance/stages/download/download_sec_filings.py @@ -20,6 +20,7 @@ import yaml from nvflow.core import BaseStage, StageRegistry, console +from nvflow.lib.cli_cmd import build_python_cmd @StageRegistry.register(recipe="finance", workflow="download-sec", stage="sap-500") @@ -30,6 +31,20 @@ class DownloadSecFilingsStage(BaseStage): workflow = "download-sec" + def validate_config(self, config: dict[str, Any]) -> None: + """Validate that required configuration fields are present.""" + for field in ("output_dir", "sec_identity_email", "sec_identity_company"): + if field not in config: + raise ValueError(f"'{field}' is required in download_sec_filings config") + + if "config" not in config: + for field in ("tickers", "start_year", "end_year"): + if field not in config: + raise ValueError( + f"'{field}' is required in download_sec_filings config " + "(or provide 'config' pointing to a filings config file)" + ) + def execute( self, config: dict[str, Any], @@ -73,17 +88,19 @@ def execute( forms_str = " ".join(forms) if isinstance(forms, list) else forms log_dir = Path(output_dir) / "download-logs" + rendered_cmd = build_python_cmd( + "nvflow.recipes.finance.utils.download.download_sec_filings", + tickers=tickers_str, + forms=forms_str, + start_year=start_year, + end_year=end_year, + output_dir=output_dir, + sec_email=sec_identity_email, + sec_company=sec_identity_company, + ) + run_cmd( - ctx=wrap_arguments( - f"python -m nvflow.recipes.finance.utils.download.download_sec_filings " - f'--tickers "{tickers_str}" ' - f'--forms "{forms_str}" ' - f"--start_year {start_year} " - f"--end_year {end_year} " - f"--output_dir {output_dir} " - f'--sec_email "{sec_identity_email}" ' - f'--sec_company "{sec_identity_company}"' - ), + ctx=wrap_arguments(rendered_cmd), cluster=cluster, **config.get("stage_kwargs", {}), expname=expname, diff --git a/tests/test_download_sec_filings.py b/tests/test_download_sec_filings.py new file mode 100644 index 0000000..7f229aa --- /dev/null +++ b/tests/test_download_sec_filings.py @@ -0,0 +1,158 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Tests for nvflow.recipes.finance.stages.download.download_sec_filings. + +``DownloadSecFilingsStage`` does not import ``nemo_skills`` at module +level (only inside ``execute()``), so this module is importable in the +lightweight CI environment. The command-rendering tests stub +``nemo_skills.pipeline.cli`` in ``sys.modules`` so ``execute()`` can run +end-to-end without the real dependency -- they assert only on what this +stage passes to ``run_cmd``/``wrap_arguments``, not on nemo-skills' +internal behavior. +""" + +from __future__ import annotations + +import sys +import types +from pathlib import Path +from typing import Any + +import pytest +import yaml + +from nvflow.recipes.finance.stages.download.download_sec_filings import ( + DownloadSecFilingsStage, +) + + +def _base_config(**overrides: Any) -> dict[str, Any]: + config = { + "output_dir": "/data/sec", + "sec_identity_email": "test@example.com", + "sec_identity_company": "Test Co", + "tickers": ["AAPL", "MSFT"], + "start_year": 2020, + "end_year": 2023, + } + config.update(overrides) + return config + + +class TestValidateConfig: + def test_valid_direct_config_passes(self) -> None: + DownloadSecFilingsStage().validate_config(_base_config()) + + def test_valid_config_file_reference_does_not_require_tickers(self) -> None: + config = { + "output_dir": "/data/sec", + "sec_identity_email": "test@example.com", + "sec_identity_company": "Test Co", + "config": "filings.yaml", + } + DownloadSecFilingsStage().validate_config(config) + + @pytest.mark.parametrize( + "missing_field", ["output_dir", "sec_identity_email", "sec_identity_company"] + ) + def test_missing_always_required_field_raises(self, missing_field: str) -> None: + config = _base_config() + del config[missing_field] + with pytest.raises(ValueError, match=missing_field): + DownloadSecFilingsStage().validate_config(config) + + @pytest.mark.parametrize("missing_field", ["tickers", "start_year", "end_year"]) + def test_missing_field_without_config_file_raises(self, missing_field: str) -> None: + config = _base_config() + del config[missing_field] + with pytest.raises(ValueError, match=missing_field): + DownloadSecFilingsStage().validate_config(config) + + +class TestExecuteCommandRendering: + """Exercises execute() with nemo_skills.pipeline.cli stubbed out. + + Captures the ``ctx`` string passed to ``run_cmd`` (our stub + ``wrap_arguments`` is the identity function) to assert on exactly + what shell command this stage builds. + """ + + @pytest.fixture + def captured_run_cmd(self, monkeypatch: pytest.MonkeyPatch) -> dict[str, Any]: + captured: dict[str, Any] = {} + + def fake_run_cmd(**kwargs: Any) -> None: + captured.update(kwargs) + + fake_cli = types.ModuleType("nemo_skills.pipeline.cli") + fake_cli.run_cmd = fake_run_cmd # type: ignore[attr-defined] + fake_cli.wrap_arguments = lambda cmd: cmd # type: ignore[attr-defined] + + fake_pipeline = types.ModuleType("nemo_skills.pipeline") + fake_pipeline.cli = fake_cli # type: ignore[attr-defined] + + fake_nemo_skills = types.ModuleType("nemo_skills") + fake_nemo_skills.pipeline = fake_pipeline # type: ignore[attr-defined] + + monkeypatch.setitem(sys.modules, "nemo_skills", fake_nemo_skills) + monkeypatch.setitem(sys.modules, "nemo_skills.pipeline", fake_pipeline) + monkeypatch.setitem(sys.modules, "nemo_skills.pipeline.cli", fake_cli) + return captured + + def test_renders_python3_module_invocation(self, captured_run_cmd: dict[str, Any]) -> None: + DownloadSecFilingsStage().execute(_base_config(), cluster="my_cluster", expname="exp") + assert captured_run_cmd["ctx"].startswith( + "python3 -m nvflow.recipes.finance.utils.download.download_sec_filings " + ) + + def test_renders_all_expected_flags(self, captured_run_cmd: dict[str, Any]) -> None: + DownloadSecFilingsStage().execute(_base_config(), cluster="my_cluster", expname="exp") + cmd = captured_run_cmd["ctx"] + assert "--tickers 'AAPL MSFT'" in cmd + assert "--start_year 2020" in cmd + assert "--end_year 2023" in cmd + assert "--output_dir /data/sec" in cmd + assert "--sec_email test@example.com" in cmd + assert "--sec_company 'Test Co'" in cmd + + def test_shell_metacharacters_in_ticker_are_safely_quoted( + self, captured_run_cmd: dict[str, Any] + ) -> None: + malicious_config = _base_config(tickers=["AAPL; rm -rf /", "$(whoami)"]) + DownloadSecFilingsStage().execute(malicious_config, cluster="my_cluster", expname="exp") + cmd = captured_run_cmd["ctx"] + # shlex.quote wraps the whole space-joined tickers string in single + # quotes, so the shell sees one literal argument -- ``;`` and + # ``$(...)`` cannot terminate the command or trigger substitution. + assert "--tickers 'AAPL; rm -rf / $(whoami)'" in cmd + + def test_loads_tickers_from_referenced_config_file( + self, captured_run_cmd: dict[str, Any], tmp_path: Path + ) -> None: + filings_config = tmp_path / "filings.yaml" + filings_config.write_text( + yaml.dump({"tickers": ["NVDA"], "start_year": 2021, "end_year": 2022}) + ) + config = { + "output_dir": "/data/sec", + "sec_identity_email": "test@example.com", + "sec_identity_company": "Test Co", + "config": str(filings_config), + } + DownloadSecFilingsStage().execute(config, cluster="my_cluster", expname="exp") + cmd = captured_run_cmd["ctx"] + assert "--tickers NVDA" in cmd + assert "--start_year 2021" in cmd + assert "--end_year 2022" in cmd