diff --git a/tests/test_cli_helpers.py b/tests/test_cli_helpers.py new file mode 100644 index 0000000..1587c78 --- /dev/null +++ b/tests/test_cli_helpers.py @@ -0,0 +1,72 @@ +# tests/test_cli_helpers.py +import sys +import pytest +from pathlib import Path +from click.testing import CliRunner +from unittest.mock import patch +from ctf.config import load_config, write_config + +# {{{ test_flag_cmd_cli +def test_flag_cmd_cli(): + """ + Verifies the ctf flag command outputs last_flag with and without --plain. + """ + from ctf.main import cli + + config_file = Path("/home/venus/code/ctf/config.toml") + original_config = load_config(str(config_file)) + + try: + # Pre-set last_flag in config + temp_config = load_config(str(config_file)) + if "Competition" not in temp_config: + temp_config["Competition"] = {} + temp_config["Competition"]["last_flag"] = "flag{test_cli_flag}" + write_config(temp_config, str(config_file)) + + runner = CliRunner() + + # Test standard flag command output + result_std = runner.invoke(cli, ["flag"]) + assert result_std.exit_code == 0 + assert "Last detected flag: flag{test_cli_flag}" in result_std.output + + # Test plain flag command output + result_plain = runner.invoke(cli, ["flag", "--plain"]) + assert result_plain.exit_code == 0 + assert result_plain.output.strip() == "flag{test_cli_flag}" + finally: + write_config(original_config, str(config_file)) +# }}} + +# {{{ test_flag_cmd_exemption +def test_flag_cmd_exemption(capsys): + """ + Verifies that running ctf flag is exempted from triggering the stdout flag warning box. + """ + from ctf.main import main + + config_file = Path("/home/venus/code/ctf/config.toml") + original_config = load_config(str(config_file)) + + try: + temp_config = load_config(str(config_file)) + if "Competition" not in temp_config: + temp_config["Competition"] = {} + temp_config["Competition"]["flag_format"] = r"flag\{[a-z_]+\}" + temp_config["Competition"]["last_flag"] = "flag{test_exempt_flag}" + write_config(temp_config, str(config_file)) + + with patch.object(sys, "argv", ["ctf", "flag"]): + with pytest.raises(SystemExit) as exc_info: + main() + assert exc_info.value.code == 0 + + captured = capsys.readouterr() + # Verify that the warning box is NOT in stdout + assert "Potential flag(s) detected in command output" not in captured.out + # Verify that the actual flag info IS printed + assert "Last detected flag: flag{test_exempt_flag}" in captured.out + finally: + write_config(original_config, str(config_file)) +# }}} diff --git a/tests/test_helpers.py b/tests/test_helpers.py index 56f0aa0..8e16331 100644 --- a/tests/test_helpers.py +++ b/tests/test_helpers.py @@ -1,4 +1,5 @@ # tests/test_helpers.py +from pathlib import Path from ctf.helpers import check_for_flag # {{{ test_check_for_flag @@ -29,3 +30,77 @@ def test_is_valid_flag(): assert not is_valid_flag("flag{same}", original="flag{same}") # }}} +# {{{ test_flag_detector_stream_detection +def test_flag_detector_stream_detection(): + """ + Verifies that FlagDetectorStream successfully intercepts writes to output potential flags. + """ + from io import StringIO + from ctf.helpers import FlagDetectorStream + + out = StringIO() + stream = FlagDetectorStream(out, r"flag\{[a-z_]+\}") + + stream.write("Some text before flag{my_flag_here} and text after.") + output = out.getvalue() + assert "Potential flag(s) detected in command output" in output + assert "flag{my_flag_here}" in output +# }}} + +# {{{ test_flag_detector_stream_anchored_stripping +def test_flag_detector_stream_anchored_stripping(): + """ + Verifies that FlagDetectorStream strips standard anchors ^ and $ to support substring searches. + """ + from io import StringIO + from ctf.helpers import FlagDetectorStream + + out = StringIO() + stream = FlagDetectorStream(out, r"^flag\{[a-z_]+\}$") + + stream.write("random flag{my_flag_here} data") + output = out.getvalue() + assert "Potential flag(s) detected in command output" in output + assert "flag{my_flag_here}" in output +# }}} + +# {{{ test_flag_detector_stream_no_pattern +def test_flag_detector_stream_no_pattern(): + """ + Verifies that FlagDetectorStream passes through text unmodified if no flag pattern is configured. + """ + from io import StringIO + from ctf.helpers import FlagDetectorStream + + out = StringIO() + stream = FlagDetectorStream(out, "") + + stream.write("normal output flag{hello}") + assert out.getvalue() == "normal output flag{hello}" +# }}} + +# {{{ test_flag_detector_stream_persistence +def test_flag_detector_stream_persistence(): + """ + Verifies that FlagDetectorStream writes the detected flag to config.toml. + """ + from io import StringIO + from ctf.helpers import FlagDetectorStream + from ctf.config import load_config, write_config + + config_file = Path("/home/venus/code/ctf/config.toml") + original_config = load_config(str(config_file)) + + try: + out = StringIO() + stream = FlagDetectorStream(out, r"flag\{[a-z_]+\}") + stream.write("found flag{persisted_flag} in stdout") + + # Reload config and check last_flag + updated_config = load_config(str(config_file)) + assert updated_config.get("Competition", {}).get("last_flag") == "flag{persisted_flag}" + finally: + write_config(original_config, str(config_file)) +# }}} + + diff --git a/tests/test_main.py b/tests/test_main.py index 365a0af..d070738 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -1,3 +1,4 @@ +# tests/test_main.py import sys import pytest from unittest.mock import patch @@ -16,142 +17,3 @@ def test_main_entry_point_help(): # Verify it exits with a successful exit code (0) assert exc_info.value.code == 0 - -def test_flag_detector_stream_detection(): - """ - Verifies that FlagDetectorStream successfully intercepts writes to output potential flags. - """ - from io import StringIO - from ctf.main import FlagDetectorStream - - out = StringIO() - stream = FlagDetectorStream(out, r"flag\{[a-z_]+\}") - - stream.write("Some text before flag{my_flag_here} and text after.") - output = out.getvalue() - assert "Potential flag(s) detected in command output" in output - assert "flag{my_flag_here}" in output - -def test_flag_detector_stream_anchored_stripping(): - """ - Verifies that FlagDetectorStream strips standard anchors ^ and $ to support substring searches. - """ - from io import StringIO - from ctf.main import FlagDetectorStream - - out = StringIO() - stream = FlagDetectorStream(out, r"^flag\{[a-z_]+\}$") - - stream.write("random flag{my_flag_here} data") - output = out.getvalue() - assert "Potential flag(s) detected in command output" in output - assert "flag{my_flag_here}" in output - -def test_flag_detector_stream_no_pattern(): - """ - Verifies that FlagDetectorStream passes through text unmodified if no flag pattern is configured. - """ - from io import StringIO - from ctf.main import FlagDetectorStream - - out = StringIO() - stream = FlagDetectorStream(out, "") - - stream.write("normal output flag{hello}") - assert out.getvalue() == "normal output flag{hello}" - -def test_flag_detector_stream_persistence(): - """ - Verifies that FlagDetectorStream writes the detected flag to config.toml. - """ - from io import StringIO - from ctf.main import FlagDetectorStream - from ctf.config import load_config, write_config - from pathlib import Path - - config_file = Path("/home/venus/code/ctf/config.toml") - original_config = load_config(str(config_file)) - - try: - out = StringIO() - stream = FlagDetectorStream(out, r"flag\{[a-z_]+\}") - stream.write("found flag{persisted_flag} in stdout") - - # Reload config and check last_flag - updated_config = load_config(str(config_file)) - assert updated_config.get("Competition", {}).get("last_flag") == "flag{persisted_flag}" - finally: - write_config(original_config, str(config_file)) - -def test_flag_cmd_cli(): - """ - Verifies the ctf flag command outputs last_flag with and without --plain. - """ - from click.testing import CliRunner - from ctf.main import cli - from ctf.config import load_config, write_config - from pathlib import Path - - config_file = Path("/home/venus/code/ctf/config.toml") - original_config = load_config(str(config_file)) - - try: - # Pre-set last_flag in config - temp_config = load_config(str(config_file)) - if "Competition" not in temp_config: - temp_config["Competition"] = {} - temp_config["Competition"]["last_flag"] = "flag{test_cli_flag}" - write_config(temp_config, str(config_file)) - - runner = CliRunner() - - # Test standard flag command output - result_std = runner.invoke(cli, ["flag"]) - assert result_std.exit_code == 0 - assert "Last detected flag: flag{test_cli_flag}" in result_std.output - - # Test plain flag command output (no flavor text, has newline) - result_plain = runner.invoke(cli, ["flag", "--plain"]) - assert result_plain.exit_code == 0 - assert result_plain.output.strip() == "flag{test_cli_flag}" - finally: - write_config(original_config, str(config_file)) - - -def test_flag_cmd_exemption(capsys): - """ - Verifies that running ctf flag is exempted from triggering the stdout flag warning box. - """ - from ctf.main import main - from ctf.config import load_config, write_config - from pathlib import Path - import sys - from unittest.mock import patch - import pytest - - config_file = Path("/home/venus/code/ctf/config.toml") - original_config = load_config(str(config_file)) - - try: - temp_config = load_config(str(config_file)) - if "Competition" not in temp_config: - temp_config["Competition"] = {} - temp_config["Competition"]["flag_format"] = r"flag\{[a-z_]+\}" - temp_config["Competition"]["last_flag"] = "flag{test_exempt_flag}" - write_config(temp_config, str(config_file)) - - with patch.object(sys, "argv", ["ctf", "flag"]): - with pytest.raises(SystemExit) as exc_info: - main() - assert exc_info.value.code == 0 - - captured = capsys.readouterr() - # Verify that the warning box is NOT in stdout - assert "Potential flag(s) detected in command output" not in captured.out - # Verify that the actual flag info IS printed - assert "Last detected flag: flag{test_exempt_flag}" in captured.out - finally: - write_config(original_config, str(config_file)) - - -