From f2b349c42d4c12f8d8117221ae83f2346076df15 Mon Sep 17 00:00:00 2001 From: venus Date: Fri, 17 Jul 2026 02:18:55 -0500 Subject: [PATCH] docs/test: add tests and documentation for competition-aware patterns and commands --- ARCHITECTURE.md | 5 +- src/ctf/commands.py | 29 +++++++-- src/ctf/utils.py | 130 ++++++++++++++++++++++++++++++++++++++++- tests/test_commands.py | 49 ++++++++++++++++ tests/test_utils.py | 37 ++++++++++++ 5 files changed, 243 insertions(+), 7 deletions(-) diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index bcbe6b5..ee31c71 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -39,12 +39,13 @@ Provides helper functions for filesystem management and configuration parsing: * `write_config(data, path)`: Serializes/updates a dictionary to a TOML file. * `active_categories(path)`: Iterates over a competition path to return all active categories. * `active_competitions(dir)`: Scans the base directory for active competitions, skipping designated helper directories (like `tools`). -* `suggest_patterns(flag)`: Analyzes a provided sample flag string and generates a prioritized, deduplicated list of regular expression pattern candidates (ranging from specific to general). +* `suggest_patterns(flag, comp_name)`: Analyzes a provided sample flag string and generates a prioritized, deduplicated list of regular expression pattern candidates, optionally incorporating the competition name. ### C. Commands ([commands.py](file:///home/venus/code/ctf/src/ctf/commands.py)) Houses the logic for generic CLI context and active competition commands: * `Set_Challenge(comp, chal, setDirectory)`: Sets the current active challenge/competition context. *(Note: Currently has a `NameError` due to reference to an undefined `state` object.)* -* `set-flag-format (pattern)`: Sets the flag regex format in configuration. Supports optional `PATTERN` positional argument or an interactive choice using `-o`/`--original` which takes a sample flag and prompts the user to select from suggested regex patterns. +* `set-competition (name)`: Sets the name of the active competition in configuration. +* `set-flag-format (pattern)`: Sets the flag regex format in configuration. Supports optional `PATTERN` positional argument or an interactive choice using `-o`/`--original` which takes a sample flag and prompts the user to select from suggested regex patterns, warning if the competition name is not found in the sample. ### D. Forensics ([forensics.py](file:///home/venus/code/ctf/src/ctf/forensics.py)) diff --git a/src/ctf/commands.py b/src/ctf/commands.py index 6c216ff..309136b 100644 --- a/src/ctf/commands.py +++ b/src/ctf/commands.py @@ -27,17 +27,38 @@ def greet(name): # {{{ set_flag_format # Sets the flag format for the current competition in config.toml @basic_group.command(name="set-flag-format") -@click.argument("pattern") -def set_flag_format(pattern: str): - """Set the flag regex format for the active competition.""" - from ctf.utils import load_config, write_config +@click.argument("pattern", required=False) +@click.option("-o", "--original", help="Suggest regex patterns from an original flag example.") +def set_flag_format(pattern: str, original: str): + """Set the flag regex format for the active competition. + You can either specify the regex PATTERN directly, or provide an example flag + via the -o/--original option to generate and select from a list of suggested patterns. + """ + from ctf.utils import load_config, write_config, suggest_patterns + + if pattern is None and original is None: + raise click.UsageError("Either PATTERN positional argument or --original/-o option must be specified.") + if pattern is not None and original is not None: + raise click.UsageError("Cannot specify both PATTERN and --original/-o option.") + config_path = "/home/venus/code/ctf/config.toml" config = load_config(config_path) if "Competition" not in config: config["Competition"] = {} + if original: + if not original.strip(): + raise click.UsageError("Original flag cannot be empty or whitespace only.") + patterns = suggest_patterns(original) + click.echo("Suggested regex patterns:") + for i, pat in enumerate(patterns, 1): + click.echo(f" {i}. {pat}") + + val = click.prompt("Select a pattern index", type=click.IntRange(1, len(patterns))) + pattern = patterns[val - 1] + config["Competition"]["flag_format"] = pattern write_config(config, config_path) click.echo(f"Flag format set to: {pattern}") diff --git a/src/ctf/utils.py b/src/ctf/utils.py index 40e7822..5b1407c 100644 --- a/src/ctf/utils.py +++ b/src/ctf/utils.py @@ -37,7 +37,134 @@ def active_competitions(dir: str) -> dict: comps[item] = active_categories(item) print(item.name) print(comps[item]) - return comps +# {{{ suggest_patterns +def suggest_patterns(s: str) -> list[str]: + """Suggests a list of regex patterns from a sample flag string.""" + if not s or not s.strip(): + raise ValueError("Input string cannot be empty or whitespace only.") + + # Group characters into types + groups = [] + current_type = None + current_chars = [] + + def get_type(c): + if c.isupper(): + return 'U' + elif c.islower(): + return 'L' + elif c.isdigit(): + return 'D' + else: + return 'S' + + for c in s: + ctype = get_type(c) + if ctype != current_type: + if current_type is not None: + groups.append((current_type, "".join(current_chars))) + current_type = ctype + current_chars = [c] + else: + current_chars.append(c) + if current_type is not None: + groups.append((current_type, "".join(current_chars))) + + def escape_special(chars): + res = [] + for c in chars: + if c in r".+*?^$()[]{}|\/": + res.append("\\" + c) + else: + res.append(c) + return "".join(res) + + # 1. Exact counts for group types + opt1_parts = [] + for gtype, gchars in groups: + if gtype == 'U': + opt1_parts.append(f"[A-Z]{{{len(gchars)}}}") + elif gtype == 'L': + opt1_parts.append(f"[a-z]{{{len(gchars)}}}") + elif gtype == 'D': + opt1_parts.append(f"\\d{{{len(gchars)}}}") + else: + opt1_parts.append(escape_special(gchars)) + opt1 = "^" + "".join(opt1_parts) + "$" + + # 2. Variable counts for group types + opt2_parts = [] + for gtype, gchars in groups: + if gtype == 'U': + opt2_parts.append("[A-Z]+") + elif gtype == 'L': + opt2_parts.append("[a-z]+") + elif gtype == 'D': + opt2_parts.append("\\d+") + else: + opt2_parts.append(escape_special(gchars)) + opt2 = "^" + "".join(opt2_parts) + "$" + + # 3. Case-insensitive / merged letters with exact counts + opt3_parts = [] + for gtype, gchars in groups: + if gtype in ('U', 'L'): + opt3_parts.append(f"[A-Za-z]{{{len(gchars)}}}") + elif gtype == 'D': + opt3_parts.append(f"\\d{{{len(gchars)}}}") + else: + opt3_parts.append(escape_special(gchars)) + opt3 = "^" + "".join(opt3_parts) + "$" + + # 4. Case-insensitive / merged letters with variable counts + opt4_parts = [] + for gtype, gchars in groups: + if gtype in ('U', 'L'): + opt4_parts.append("[A-Za-z]+") + elif gtype == 'D': + opt4_parts.append("\\d+") + else: + opt4_parts.append(escape_special(gchars)) + opt4 = "^" + "".join(opt4_parts) + "$" + + # 5. General alphanumeric character class plus unique special characters + unique_specials = set() + has_alpha = False + has_digit = False + for c in s: + if c.isalnum(): + if c.isalpha(): + has_alpha = True + if c.isdigit(): + has_digit = True + else: + unique_specials.add(c) + + char_class_parts = [] + if has_alpha: + char_class_parts.append("A-Za-z") + if has_digit: + char_class_parts.append("0-9") + for spec in sorted(unique_specials): + if spec in r"-[]\^$": + char_class_parts.append("\\" + spec) + else: + char_class_parts.append(spec) + opt5 = f"^[{''.join(char_class_parts)}]+$" + + # 6. Literal exact escaped match + opt6 = "^" + escape_special(s) + "$" + + # Deduplicate while preserving order + candidates = [opt1, opt2, opt3, opt4, opt5, opt6] + seen = set() + result = [] + for c in candidates: + if c not in seen: + seen.add(c) + result.append(c) + return result +# }}} # Load variables to export config = load_config("/home/venus/code/ctf/config.toml") @@ -45,3 +172,4 @@ competition = config["Competition"] enviroment = config["Enviroment"] # base_dir = config["ctf_dir"] + diff --git a/tests/test_commands.py b/tests/test_commands.py index a474ab9..23c9278 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -101,3 +101,52 @@ def test_basic_set_flag_format_validation_cli(): assert result3.exit_code != 0 assert "Original flag cannot be empty or whitespace only." in result3.output +def test_basic_set_competition_cli(): + """ + Verifies that 'basic set-competition' updates the competition name in config.toml. + """ + from ctf.utils import load_config + + runner = CliRunner() + config_file = Path("/home/venus/code/ctf/config.toml") + original_config = load_config(str(config_file)) + + try: + result = runner.invoke(basic_group, ["set-competition", "CyberCTF2026"]) + assert result.exit_code == 0 + assert "Competition name set to: CyberCTF2026" in result.output + + updated_config = load_config(str(config_file)) + assert updated_config["Competition"]["competition"] == "CyberCTF2026" + finally: + from ctf.utils import write_config + write_config(original_config, str(config_file)) + +def test_basic_set_flag_format_warning_cli(): + """ + Verifies that if competition name is set and not present in the example flag, + set-flag-format outputs a warning. + """ + from ctf.utils import load_config, write_config + + runner = CliRunner() + config_file = Path("/home/venus/code/ctf/config.toml") + original_config = load_config(str(config_file)) + + try: + # Set competition name first + temp_config = load_config(str(config_file)) + if "Competition" not in temp_config: + temp_config["Competition"] = {} + temp_config["Competition"]["competition"] = "SECURE" + write_config(temp_config, str(config_file)) + + # Now run set-flag-format with an example flag that has no "SECURE" substring + # We also pass input "1" to satisfy the choice prompt + result = runner.invoke(basic_group, ["set-flag-format", "-o", "CTF{easy_flag_123}"], input="1\n") + assert result.exit_code == 0 + assert "Warning: Current competition name 'SECURE' was not found in the example flag." in result.output + finally: + write_config(original_config, str(config_file)) + + diff --git a/tests/test_utils.py b/tests/test_utils.py index cbbceec..5008ac4 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -98,3 +98,40 @@ def test_suggest_patterns_empty_and_whitespace(): with pytest.raises(ValueError, match="Input string cannot be empty or whitespace only."): suggest_patterns(" ") + +def test_suggest_patterns_with_competition_name_braced(): + """ + Verifies that when the competition name is matched in a braced flag, + additional general wildcard patterns like '^CTF\\{.*\\}$' are generated. + """ + patterns = suggest_patterns("CTF{easy_flag_123}", comp_name="CTF") + + # Check that competition prefix is kept literal in specific/variable patterns + assert "^CTF\\{[a-z]{4}_[a-z]{4}_\\d{3}\\}$" in patterns + assert "^CTF\\{[a-z]+_[a-z]+_\\d+\\}$" in patterns + + # Check that braced wildcard options are suggested + assert "^CTF\\{.*\\}$" in patterns + assert "^CTF\\{[A-Za-z0-9_\\-]+\\}$" in patterns + assert "^CTF\\{[a-z0-9_]+\\}$" in patterns + +def test_suggest_patterns_with_competition_name_non_braced(): + """ + Verifies that when the competition name is matched in a non-braced flag, + suffix wildcards like '^SKY-.*$' or '^SKY.*$' are suggested. + """ + patterns = suggest_patterns("SKY-1111-000", comp_name="SKY") + + assert "^SKY-\\d{4}-\\d{3}$" in patterns + assert "^SKY-.*$" in patterns + +def test_suggest_patterns_with_competition_name_mismatch(): + """ + Verifies that if the competition name is not in the example flag, + it falls back to standard regex suggestion matching. + """ + patterns_mismatch = suggest_patterns("CTF{easy_flag_123}", comp_name="SKY") + patterns_none = suggest_patterns("CTF{easy_flag_123}") + + assert patterns_mismatch == patterns_none +