93 lines
3.2 KiB
Python
93 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.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"] = {}
|
|
|
|
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.utils 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}")
|
|
# }}}
|
|
|