Files
CTF-tool/src/ctf/commands.py

94 lines
3.2 KiB
Python

# functions for commands needed
# src/commands.py
# vim foldmethod=marker
import click
from pathlib import Path
# {{{ basic_group
# This defines a group with name basic which will nest other comands to be imported in the main loop
@click.group(name="basic")
def basic_group(): pass
# }}}
# {{{ test
# Adds a simple commmand to be run with `ctf basic test`
@basic_group.command(name="test")
def test():
print("hello from test")
# }}}
# {{{ greet
# A simple command with a positional(required) argument name
@basic_group.command(name="greet")
@click.argument('name')
def greet(name):
print(f"hello {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", 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.config import load_config, write_config
from ctf.utils import 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"] = {}
comp_name = config["Competition"].get("competition", "")
if original:
if not original.strip():
raise click.UsageError("Original flag cannot be empty or whitespace only.")
if comp_name and comp_name.lower() not in original.lower():
click.echo(f"Warning: Current competition name '{comp_name}' was not found in the example flag.")
patterns = suggest_patterns(original, comp_name=comp_name)
click.echo(f"Suggested regex patterns for competition {comp_name}:")
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}")
# }}}
# {{{ set_competition
# Sets the competition name in config.toml
@basic_group.command(name="set-competition")
@click.argument("name")
def set_competition(name: str):
"""Set the name of the active competition."""
from ctf.config import load_config, write_config
config_path = "/home/venus/code/ctf/config.toml"
config = load_config(config_path)
if "Competition" not in config:
config["Competition"] = {}
config["Competition"]["competition"] = name
write_config(config, config_path)
click.echo(f"Competition name set to: {name}")
# }}}