diff --git a/tests/test_main.py b/tests/test_main.py index 4f665df..4a68d53 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -60,3 +60,61 @@ def test_flag_detector_stream_no_pattern(): 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.utils 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.utils 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, no newline) + result_plain = runner.invoke(cli, ["flag", "--plain"]) + assert result_plain.exit_code == 0 + assert result_plain.output == "flag{test_cli_flag}" + finally: + write_config(original_config, str(config_file)) + +