Skip to content

Commit d823175

Browse files
Anilreddy2309claude
andcommitted
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 <noreply@anthropic.com> Signed-off-by: Anil Balireddy <anilbalireddi@gmail.com>
1 parent 9f34b6b commit d823175

2 files changed

Lines changed: 185 additions & 10 deletions

File tree

nvflow/recipes/finance/stages/download/download_sec_filings.py

Lines changed: 27 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import yaml
2121

2222
from nvflow.core import BaseStage, StageRegistry, console
23+
from nvflow.lib.cli_cmd import build_python_cmd
2324

2425

2526
@StageRegistry.register(recipe="finance", workflow="download-sec", stage="sap-500")
@@ -30,6 +31,20 @@ class DownloadSecFilingsStage(BaseStage):
3031

3132
workflow = "download-sec"
3233

34+
def validate_config(self, config: dict[str, Any]) -> None:
35+
"""Validate that required configuration fields are present."""
36+
for field in ("output_dir", "sec_identity_email", "sec_identity_company"):
37+
if field not in config:
38+
raise ValueError(f"'{field}' is required in download_sec_filings config")
39+
40+
if "config" not in config:
41+
for field in ("tickers", "start_year", "end_year"):
42+
if field not in config:
43+
raise ValueError(
44+
f"'{field}' is required in download_sec_filings config "
45+
"(or provide 'config' pointing to a filings config file)"
46+
)
47+
3348
def execute(
3449
self,
3550
config: dict[str, Any],
@@ -73,17 +88,19 @@ def execute(
7388
forms_str = " ".join(forms) if isinstance(forms, list) else forms
7489
log_dir = Path(output_dir) / "download-logs"
7590

91+
rendered_cmd = build_python_cmd(
92+
"nvflow.recipes.finance.utils.download.download_sec_filings",
93+
tickers=tickers_str,
94+
forms=forms_str,
95+
start_year=start_year,
96+
end_year=end_year,
97+
output_dir=output_dir,
98+
sec_email=sec_identity_email,
99+
sec_company=sec_identity_company,
100+
)
101+
76102
run_cmd(
77-
ctx=wrap_arguments(
78-
f"python -m nvflow.recipes.finance.utils.download.download_sec_filings "
79-
f'--tickers "{tickers_str}" '
80-
f'--forms "{forms_str}" '
81-
f"--start_year {start_year} "
82-
f"--end_year {end_year} "
83-
f"--output_dir {output_dir} "
84-
f'--sec_email "{sec_identity_email}" '
85-
f'--sec_company "{sec_identity_company}"'
86-
),
103+
ctx=wrap_arguments(rendered_cmd),
87104
cluster=cluster,
88105
**config.get("stage_kwargs", {}),
89106
expname=expname,

tests/test_download_sec_filings.py

Lines changed: 158 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,158 @@
1+
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
#
15+
"""Tests for nvflow.recipes.finance.stages.download.download_sec_filings.
16+
17+
``DownloadSecFilingsStage`` does not import ``nemo_skills`` at module
18+
level (only inside ``execute()``), so this module is importable in the
19+
lightweight CI environment. The command-rendering tests stub
20+
``nemo_skills.pipeline.cli`` in ``sys.modules`` so ``execute()`` can run
21+
end-to-end without the real dependency -- they assert only on what this
22+
stage passes to ``run_cmd``/``wrap_arguments``, not on nemo-skills'
23+
internal behavior.
24+
"""
25+
26+
from __future__ import annotations
27+
28+
import sys
29+
import types
30+
from pathlib import Path
31+
from typing import Any
32+
33+
import pytest
34+
import yaml
35+
36+
from nvflow.recipes.finance.stages.download.download_sec_filings import (
37+
DownloadSecFilingsStage,
38+
)
39+
40+
41+
def _base_config(**overrides: Any) -> dict[str, Any]:
42+
config = {
43+
"output_dir": "/data/sec",
44+
"sec_identity_email": "test@example.com",
45+
"sec_identity_company": "Test Co",
46+
"tickers": ["AAPL", "MSFT"],
47+
"start_year": 2020,
48+
"end_year": 2023,
49+
}
50+
config.update(overrides)
51+
return config
52+
53+
54+
class TestValidateConfig:
55+
def test_valid_direct_config_passes(self) -> None:
56+
DownloadSecFilingsStage().validate_config(_base_config())
57+
58+
def test_valid_config_file_reference_does_not_require_tickers(self) -> None:
59+
config = {
60+
"output_dir": "/data/sec",
61+
"sec_identity_email": "test@example.com",
62+
"sec_identity_company": "Test Co",
63+
"config": "filings.yaml",
64+
}
65+
DownloadSecFilingsStage().validate_config(config)
66+
67+
@pytest.mark.parametrize(
68+
"missing_field", ["output_dir", "sec_identity_email", "sec_identity_company"]
69+
)
70+
def test_missing_always_required_field_raises(self, missing_field: str) -> None:
71+
config = _base_config()
72+
del config[missing_field]
73+
with pytest.raises(ValueError, match=missing_field):
74+
DownloadSecFilingsStage().validate_config(config)
75+
76+
@pytest.mark.parametrize("missing_field", ["tickers", "start_year", "end_year"])
77+
def test_missing_field_without_config_file_raises(self, missing_field: str) -> None:
78+
config = _base_config()
79+
del config[missing_field]
80+
with pytest.raises(ValueError, match=missing_field):
81+
DownloadSecFilingsStage().validate_config(config)
82+
83+
84+
class TestExecuteCommandRendering:
85+
"""Exercises execute() with nemo_skills.pipeline.cli stubbed out.
86+
87+
Captures the ``ctx`` string passed to ``run_cmd`` (our stub
88+
``wrap_arguments`` is the identity function) to assert on exactly
89+
what shell command this stage builds.
90+
"""
91+
92+
@pytest.fixture
93+
def captured_run_cmd(self, monkeypatch: pytest.MonkeyPatch) -> dict[str, Any]:
94+
captured: dict[str, Any] = {}
95+
96+
def fake_run_cmd(**kwargs: Any) -> None:
97+
captured.update(kwargs)
98+
99+
fake_cli = types.ModuleType("nemo_skills.pipeline.cli")
100+
fake_cli.run_cmd = fake_run_cmd # type: ignore[attr-defined]
101+
fake_cli.wrap_arguments = lambda cmd: cmd # type: ignore[attr-defined]
102+
103+
fake_pipeline = types.ModuleType("nemo_skills.pipeline")
104+
fake_pipeline.cli = fake_cli # type: ignore[attr-defined]
105+
106+
fake_nemo_skills = types.ModuleType("nemo_skills")
107+
fake_nemo_skills.pipeline = fake_pipeline # type: ignore[attr-defined]
108+
109+
monkeypatch.setitem(sys.modules, "nemo_skills", fake_nemo_skills)
110+
monkeypatch.setitem(sys.modules, "nemo_skills.pipeline", fake_pipeline)
111+
monkeypatch.setitem(sys.modules, "nemo_skills.pipeline.cli", fake_cli)
112+
return captured
113+
114+
def test_renders_python3_module_invocation(self, captured_run_cmd: dict[str, Any]) -> None:
115+
DownloadSecFilingsStage().execute(_base_config(), cluster="my_cluster", expname="exp")
116+
assert captured_run_cmd["ctx"].startswith(
117+
"python3 -m nvflow.recipes.finance.utils.download.download_sec_filings "
118+
)
119+
120+
def test_renders_all_expected_flags(self, captured_run_cmd: dict[str, Any]) -> None:
121+
DownloadSecFilingsStage().execute(_base_config(), cluster="my_cluster", expname="exp")
122+
cmd = captured_run_cmd["ctx"]
123+
assert "--tickers 'AAPL MSFT'" in cmd
124+
assert "--start_year 2020" in cmd
125+
assert "--end_year 2023" in cmd
126+
assert "--output_dir /data/sec" in cmd
127+
assert "--sec_email test@example.com" in cmd
128+
assert "--sec_company 'Test Co'" in cmd
129+
130+
def test_shell_metacharacters_in_ticker_are_safely_quoted(
131+
self, captured_run_cmd: dict[str, Any]
132+
) -> None:
133+
malicious_config = _base_config(tickers=["AAPL; rm -rf /", "$(whoami)"])
134+
DownloadSecFilingsStage().execute(malicious_config, cluster="my_cluster", expname="exp")
135+
cmd = captured_run_cmd["ctx"]
136+
# shlex.quote wraps the whole space-joined tickers string in single
137+
# quotes, so the shell sees one literal argument -- ``;`` and
138+
# ``$(...)`` cannot terminate the command or trigger substitution.
139+
assert "--tickers 'AAPL; rm -rf / $(whoami)'" in cmd
140+
141+
def test_loads_tickers_from_referenced_config_file(
142+
self, captured_run_cmd: dict[str, Any], tmp_path: Path
143+
) -> None:
144+
filings_config = tmp_path / "filings.yaml"
145+
filings_config.write_text(
146+
yaml.dump({"tickers": ["NVDA"], "start_year": 2021, "end_year": 2022})
147+
)
148+
config = {
149+
"output_dir": "/data/sec",
150+
"sec_identity_email": "test@example.com",
151+
"sec_identity_company": "Test Co",
152+
"config": str(filings_config),
153+
}
154+
DownloadSecFilingsStage().execute(config, cluster="my_cluster", expname="exp")
155+
cmd = captured_run_cmd["ctx"]
156+
assert "--tickers NVDA" in cmd
157+
assert "--start_year 2021" in cmd
158+
assert "--end_year 2022" in cmd

0 commit comments

Comments
 (0)