venus/CTF-tool
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

docs/test: add tests and documentation for competition-aware patterns and commands

Browse Source
This commit is contained in:
venus
2026-07-17 02:18:55 -05:00
parent 9e91f8ef0f
commit f2b349c42d
5 changed files with 243 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
ARCHITECTURE.md
View File

@@ -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))

29
src/ctf/commands.py
View File

@@ -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}")

130
src/ctf/utils.py
View File

@@ -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"]

49
tests/test_commands.py
View File

@@ -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))

37
tests/test_utils.py
View File

@@ -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