128 lines
4.4 KiB
Python
128 lines
4.4 KiB
Python
# 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))
|
|
# }}}
|
|
|
|
# {{{ test_flag_cmd_list_format
|
|
def test_flag_cmd_list_format():
|
|
"""
|
|
Verifies that ctf flag --list / -l prints the current flag format.
|
|
"""
|
|
from ctf.main import cli
|
|
|
|
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"] = "TEST_FORMAT{[a-z]+}"
|
|
write_config(temp_config, str(config_file))
|
|
|
|
runner = CliRunner()
|
|
|
|
# Test standard output
|
|
result = runner.invoke(cli, ["flag", "-l"])
|
|
assert result.exit_code == 0
|
|
assert "Current flag format: TEST_FORMAT{[a-z]+}" in result.output
|
|
|
|
# Test plain output
|
|
result_plain = runner.invoke(cli, ["flag", "-l", "-p"])
|
|
assert result_plain.exit_code == 0
|
|
assert result_plain.output.strip() == "TEST_FORMAT{[a-z]+}"
|
|
finally:
|
|
write_config(original_config, str(config_file))
|
|
# }}}
|
|
|
|
# {{{ test_flag_cmd_set_format
|
|
def test_flag_cmd_set_format():
|
|
"""
|
|
Verifies that ctf flag --set / -s updates the flag format in config.toml.
|
|
"""
|
|
from ctf.main import cli
|
|
|
|
config_file = Path("/home/venus/code/ctf/config.toml")
|
|
original_config = load_config(str(config_file))
|
|
|
|
try:
|
|
runner = CliRunner()
|
|
result = runner.invoke(cli, ["flag", "-s", "NEW_FLAG_FORMAT{[0-9]+}"])
|
|
assert result.exit_code == 0
|
|
assert "Flag format set to: NEW_FLAG_FORMAT{[0-9]+}" in result.output
|
|
|
|
updated_config = load_config(str(config_file))
|
|
assert updated_config["Competition"]["flag_format"] == "NEW_FLAG_FORMAT{[0-9]+}"
|
|
finally:
|
|
write_config(original_config, str(config_file))
|
|
# }}}
|
|
|