From c069e512639cf5b6847a60e34d71e940064a2d2d Mon Sep 17 00:00:00 2001 From: venus Date: Fri, 17 Jul 2026 01:55:44 -0500 Subject: [PATCH] Standardize matching test files structure, add test_main.py and test_commands.py, and update GEMINI.md --- GEMINI.md | 1 + src/ctf/commands.py | 37 +++++++++++++++++----- src/ctf/forensics.py | 68 +++++++++++++++++++++++++++++++++++++++- src/ctf/main.py | 11 ++++--- tests/env/clean_data.bin | 1 + tests/env/flag_data.bin | 2 +- tests/test_commands.py | 49 +++++++++++++++++++++++++++++ tests/test_forensics.py | 26 --------------- tests/test_main.py | 18 +++++++++++ 9 files changed, 173 insertions(+), 40 deletions(-) create mode 100644 tests/env/clean_data.bin create mode 100644 tests/test_commands.py create mode 100644 tests/test_main.py diff --git a/GEMINI.md b/GEMINI.md index 70c88dc..9e18bff 100644 --- a/GEMINI.md +++ b/GEMINI.md @@ -11,6 +11,7 @@ The primary goal is to use the context of CTF challenges (forensics, crypto, web 3. **Summarization:** When directed, summarize documentation (local or web-based) to aid in understanding CTF tools and Python modules. 4. **Educational Context:** Use the existing scripts and tools in this repository (like `psk_crack.py`, `exploit.sh`, or the `tools/` directory) as examples when explaining technical concepts. 5. **Vim Folding Markers:** Wrap all classes, command groups, subcommands, and functions inside project implementation files in standard Vim/Neovim folding syntax markers (`# {{{ ` and `# }}}`). +6. **Matching Test Files:** Every Python implementation file inside `src/` must have a corresponding test file under the `tests/` directory named `test_.py`. ## Interaction Workflow When discussing new implementations, features, or additions: diff --git a/src/ctf/commands.py b/src/ctf/commands.py index 87269cf..6c216ff 100644 --- a/src/ctf/commands.py +++ b/src/ctf/commands.py @@ -1,22 +1,47 @@ # functions for commands needed # src/commands.py -# from pathlib import path 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 +# }}} -# Adds a simple commmand to be run with `ctf test` per the function name -@click.command() +# {{{ 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 -@click.command() +@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") +def set_flag_format(pattern: str): + """Set the flag regex format for 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"]["flag_format"] = pattern + write_config(config, config_path) + click.echo(f"Flag format set to: {pattern}") +# }}} # def Set_Challenge(comp: str, chal: str, setDirectory: bool): # # set the current challenge and competition from input @@ -24,11 +49,9 @@ def greet(name): # state.current_comp=comp # state.comp_dir=pathlib # # TODO archive the old competitions - +# # if state.current_chal != chal: # state.current_chal=chal # print("challenge already set") # # TODO archive the old challenges # # TODO set the directory to challenge directory, with ignore option - - diff --git a/src/ctf/forensics.py b/src/ctf/forensics.py index a418610..bc69f10 100644 --- a/src/ctf/forensics.py +++ b/src/ctf/forensics.py @@ -1,6 +1,7 @@ # src/ctf/forensics.py # Library for forensic analysis +# vim foldmethod=marker import click from dataclasses import dataclass, field from pathlib import Path @@ -34,18 +35,24 @@ COMMON_SIGNATURES = { b"ID3": ("MP3 Audio", [".mp3"]), } +# {{{ forensics_group @click.group(name="forensics") -def forensics_group(): pass +def forensics_group(): + ''' A collection of forensics tools ''' + pass @forensics_group.command() def tf(): print("hello from forensics") +# }}} +# {{{ FileMetadata @dataclass class FileMetadata: filename: str size: int magic: str extension: str + detected_type: str # Task 1: POSIX Permissions permissions_octal: str @@ -69,10 +76,13 @@ class FileMetadata: # Task 6: Extended Attributes extended_attributes: Dict[str, str] = field(default_factory=dict) +# }}} +# {{{ inspect @forensics_group.command() @click.argument('path', type=click.Path(exists=True)) def inspect(path): + '''Lists all basic inode metadata about a file''' p = Path(path) stat_info = p.stat() @@ -87,6 +97,16 @@ def inspect(path): except Exception: magic = "" + detected_type = "Unknown" + try: + magic_bytes = bytes.fromhex(magic) + for signature, (type_name, exts) in COMMON_SIGNATURES.items(): + if magic_bytes.startswith(signature): + detected_type = type_name + break + except Exception: + pass + # POSIX Permissions mode = stat_info.st_mode permissions_octal = oct(stat.S_IMODE(mode)) @@ -141,6 +161,7 @@ def inspect(path): size=size, magic=magic, extension=extension, + detected_type=detected_type, permissions_octal=permissions_octal, permissions_symbolic=permissions_symbolic, owner_uid=owner_uid, @@ -166,6 +187,7 @@ def inspect(path): table.add_row("Filename", meta.filename) table.add_row("Size", f"{meta.size} bytes") table.add_row("Magic Bytes (Hex)", meta.magic) + table.add_row("Detected Type", meta.detected_type) table.add_row("Extension", meta.extension) table.add_row("Permissions", f"{meta.permissions_symbolic} ({meta.permissions_octal})") table.add_row("Owner", f"{meta.owner_username} (UID: {meta.owner_uid})") @@ -181,7 +203,9 @@ def inspect(path): console.print(table) return meta +# }}} +# {{{ list_signatures @forensics_group.command(name="signatures") def list_signatures(): """List all supported file magic signatures and expected extensions.""" @@ -200,3 +224,45 @@ def list_signatures(): table.add_row(hex_str, type_name, exts_str) console.print(table) +# }}} + +# {{{ flag_detect +@forensics_group.command(name="flag-detect") +@click.argument("filepath", type=click.Path(exists=True, dir_okay=False, path_type=Path)) +@click.option("--pattern", "-p", default=r"[a-zA-Z0-9_\-]+{[^}]+}", help="Regex pattern to search for.") +@click.option("--min-len", "-m", default=4, help="Minimum printable string length.") +def flag_detect(filepath: Path, pattern: str, min_len: int): + """Extract printable strings and search for flag patterns.""" + import re + from rich.console import Console + from rich.panel import Panel + + console = Console() + try: + with open(filepath, "rb") as f: + content = f.read() + except Exception as e: + console.print(f"[bold red]Error reading file:[/bold red] {e}") + return + + # Extract printable strings + printable_re = re.compile(rb"[a-zA-Z0-9/\-:.,_$%'\"()[\]<> ]{" + str(min_len).encode() + rb",}") + strings = printable_re.findall(content) + + # Search for flags matching pattern + flag_re = re.compile(pattern.encode("utf-8")) + flags_found = [] + for s in strings: + for match in flag_re.finditer(s): + flags_found.append(match.group().decode("utf-8", errors="ignore")) + + if flags_found: + output = "\n".join(f" [bold green]✓[/bold green] {flag}" for flag in flags_found) + console.print(Panel( + output, + title=f"[bold green]Potential Flag(s) Found ({len(flags_found)})[/bold green]", + border_style="green" + )) + else: + console.print("[bold yellow]No flag patterns found.[/bold yellow]") +# }}} diff --git a/src/ctf/main.py b/src/ctf/main.py index d411c87..9cc9968 100644 --- a/src/ctf/main.py +++ b/src/ctf/main.py @@ -1,18 +1,19 @@ # src/main.py # Parses and calls commands -import ctf.commands as commands +from ctf.commands import basic_group from ctf.forensics import forensics_group import click - +# {{{ main def main(): @click.group() def cli(): pass - cli.add_command(forensics_group()) - - + cli.add_command(forensics_group) + cli.add_command(basic_group) + cli() +# }}} if __name__ == "__main__": main() diff --git a/tests/env/clean_data.bin b/tests/env/clean_data.bin new file mode 100644 index 0000000..c9c6c1b --- /dev/null +++ b/tests/env/clean_data.bin @@ -0,0 +1 @@ +this is a completely normal text file without flags. \ No newline at end of file diff --git a/tests/env/flag_data.bin b/tests/env/flag_data.bin index 158299e..b777fef 100644 --- a/tests/env/flag_data.bin +++ b/tests/env/flag_data.bin @@ -1 +1 @@ -garbage_data_here_flag{f0rens1cs_1s_fun}more_garbage \ No newline at end of file +random_data_here_flag{found_statically_in_file}more_data \ No newline at end of file diff --git a/tests/test_commands.py b/tests/test_commands.py new file mode 100644 index 0000000..bcddb3a --- /dev/null +++ b/tests/test_commands.py @@ -0,0 +1,49 @@ +from pathlib import Path +import toml +from click.testing import CliRunner +from ctf.commands import basic_group + +TEST_ENV = Path("tests/env") + +def test_basic_test_cli(): + """ + Verifies that 'basic test' runs successfully and outputs 'hello from test'. + """ + runner = CliRunner() + result = runner.invoke(basic_group, ["test"]) + assert result.exit_code == 0 + assert "hello from test" in result.output + +def test_basic_greet_cli(): + """ + Verifies that 'basic greet' runs successfully and greets the name argument. + """ + runner = CliRunner() + result = runner.invoke(basic_group, ["greet", "Alice"]) + assert result.exit_code == 0 + assert "hello Alice" in result.output + +def test_basic_set_flag_format_cli(): + """ + Verifies that 'basic set-flag-format' CLI command writes the pattern to config.toml. + """ + from ctf.utils import load_config + + runner = CliRunner() + config_file = Path("/home/venus/code/ctf/config.toml") + + # Save the original config to restore later + original_config = load_config(str(config_file)) + + try: + result = runner.invoke(basic_group, ["set-flag-format", "TEST_FLAG{[a-z]+}"]) + assert result.exit_code == 0 + assert "Flag format set to" in result.output + + # Verify it was written to config.toml + updated_config = load_config(str(config_file)) + assert updated_config["Competition"]["flag_format"] == "TEST_FLAG{[a-z]+}" + finally: + # Restore original config + from ctf.utils import write_config + write_config(original_config, str(config_file)) diff --git a/tests/test_forensics.py b/tests/test_forensics.py index 511c02b..805f022 100644 --- a/tests/test_forensics.py +++ b/tests/test_forensics.py @@ -295,29 +295,3 @@ def test_forensics_flag_detect_cli_not_found(): result = runner.invoke(forensics_group, ["flag-detect", str(test_file)]) assert result.exit_code == 0 assert "No flag patterns found" in result.output - -def test_set_flag_format_cli(): - """ - Verifies that 'set-flag-format' CLI command writes the pattern to config.toml. - """ - from ctf.commands import set_flag_format - from ctf.utils import load_config - - runner = CliRunner() - config_file = Path("/home/venus/code/ctf/config.toml") - - # Save the original config to restore later - original_config = load_config(str(config_file)) - - try: - result = runner.invoke(set_flag_format, ["TEST_FLAG{[a-z]+}"]) - assert result.exit_code == 0 - assert "Flag format set to" in result.output - - # Verify it was written to config.toml - updated_config = load_config(str(config_file)) - assert updated_config["Competition"]["flag_format"] == "TEST_FLAG{[a-z]+}" - finally: - # Restore original config - from ctf.utils import write_config - write_config(original_config, str(config_file)) diff --git a/tests/test_main.py b/tests/test_main.py new file mode 100644 index 0000000..873ac15 --- /dev/null +++ b/tests/test_main.py @@ -0,0 +1,18 @@ +import sys +import pytest +from unittest.mock import patch +from ctf.main import main + +def test_main_entry_point_help(): + """ + Verifies that calling main() executes the Click CLI app and responds + to '--help' by listing registration groups. + """ + # Mock sys.argv to simulate running 'ctf --help' from the shell + with patch.object(sys, "argv", ["ctf", "--help"]): + # Click calls sys.exit() after displaying help, raising SystemExit + with pytest.raises(SystemExit) as exc_info: + main() + + # Verify it exits with a successful exit code (0) + assert exc_info.value.code == 0