Tool for managing CTF Challenges
from ctf_architect.core.challenge import *
from ctf_architect.core.config import *
from ctf_architect.core.mappings import *
from ctf_architect.core.stats import *
from ctf_architect.core.unzip import *

ctf_architect/__init__.py
Jus-Codin/CTF-Architect
from __future__ import annotations

from ctf_architect.cli import app
from ctf_architect.core.constants import APP_CMD_NAME


app(prog_name=APP_CMD_NAME)
ctf_architect/__main__.py
Jus-Codin/CTF-Architect
from __future__ import annotations

from ctf_architect.chall_architect.cli import app


app(prog_name="chall-architect")
ctf_architect/chall_architect/__main__.py
Jus-Codin/CTF-Architect
"""
A cli to create challenges and export them as a zip file.
"""

from __future__ import annotations

import time
import tkinter as tk
from pathlib import Path
from tkinter.filedialog import askdirectory, askopenfilename, askopenfilenames

import typer
from rich.console import Console
from rich.panel import Panel
from rich.prompt import Confirm, IntPrompt, Prompt

from ctf_architect.chall_architect.create import create_challenge
from ctf_architect.chall_architect.utils import is_valid_service_folder, get_config


try:
  import ctypes
  # Make it not blurry on windows
  ctypes.windll.shcore.SetProcessDpiAwareness(1)
except:
  pass


# This is to fix a bug on vscode's terminal where the filedialog window will not show up
_root = tk.Tk()
_root.withdraw()
_root.call('wm', 'attributes', '.', '-topmost', True)


app = typer.Typer()


console = Console()


def prompt(message: str, allow_empty: bool = False, method = Prompt.ask, **kwargs) -> str:
  """
  Prompts the user for input.
  """
  if not allow_empty:
    while True:
      res = method(message, **kwargs)
      if res != "":
        return res
      else:
        console.print("[bright_red]This field cannot be empty.[/bright_red]")
  else:
    return method(message)
  

def get_ctf_config():
  """
  Prompts the user for the ctf config.
  """
  console.rule(":gear: [bold yellow]CTF Config[/] :gear:")

  while True:

    console.print("[cyan]Please select the CTF config file.[/cyan]")
    time.sleep(1)

    config_file = askopenfilename(title="Please select the ctf config", filetypes=[("YAML", "*.yaml")])

    if config_file == "":
      # User cancelled, abort
      console.print("[bright_red]No file selected, aborting...[/bright_red]")
      return None
  
    try:
      config = get_config(config_file)
    except Exception as e:
      console.print(f"[bright_red]Error loading config file: {e}[/bright_red]")
      return None
    
    # Print config in a panel
    config_string = f"[cyan][bold]Categories:[/bold]"
    for category in config.categories:
      config_string += f"\n- {category.capitalize()}"

    config_string += f"\n\n[bold]Difficulties:[/bold]"
    for difficulty in config.difficulties:
      config_string += f"\n- {difficulty['name'].capitalize()} ({difficulty['value']})"
    config_string += "[/cyan]"

    config_name = config.name if config.name is not None else "Unnamed"

    console.print(Panel(
      config_string,
      border_style="green",
      title=f"[bright_yellow]{config_name} CTF Config[/bright_yellow]"
    ))

    if Confirm.ask("[cyan]Is this the correct config?[/]"):
      return config
    
    # Add spacing
    console.print()


def get_description():
  """
  Prompts the user for the description of the challenge.
  This should allow users to enter multiple lines, and exit when EOF is entered.
  """

  console.print(":rocket: [cyan][2/6] Please enter the challenge description. Press Ctrl+D (or Ctrl+Z on Windows) on an empty line to finish.[/]", highlight=False)

  description = []
  while True:
    try:
      line = input(">>> ")
      description.append(line)
    except EOFError:
      break
    except KeyboardInterrupt:
      return None
  
  return "\n".join(description)
  

def get_solution_files():
  """
  Prompts the user for the solution files.
  """
  console.rule(":file_folder: [bold yellow]Please select the solution files for the challenge.[/] :file_folder:")

  time.sleep(1)

  return askopenfilenames(title="Please select the solution files")


def get_flags() -> list[dict[str, str | bool]]:
  """
  Prompts the user for the flags.
  """
  console.rule(":triangular_flag: [bold yellow]Challenge Flags[/] :triangular_flag:")

  flags = []
  while True:
    regex = Confirm.ask(":triangular_flag: [cyan]Is the flag a regex flag?[/]")
    flag = prompt(":triangular_flag: [cyan]Please enter the flag[/]")
    flags.append({
      "flag": flag,
      "regex": regex
    })

    if not Confirm.ask("[cyan]Do you want to add another flag?[/]"):
      break

    # Add spacing between the flags
    console.print()

  return flags


def get_hints() -> list[dict[str, str | int | list[int]]] | None:
  """
  Prompts the user for the hints.
  """
  console.rule(":light_bulb: [bold yellow]Challenge Hints[/] :light_bulb:")

  has_hints = Confirm.ask(":light_bulb: [cyan]Does the challenge have hints?[/]")
  if has_hints:
    hints = []

    # Add spacing
    console.print()

    while True:
      hint_description = prompt(":light_bulb: [cyan]Please enter the hint description[/]")
      cost = prompt(":light_bulb: [cyan]Please enter the hint cost[/]", method=IntPrompt.ask)
      
      hint_req = None
      if len(hints) > 0:
        if Confirm.ask(":light_bulb: [cyan]Do you want to add a hint requirement?[/]"):
          # console.print the following hints in a list, with an index starting at 0
          console.print("\n[green]Hints:[/green]")
          console.print("\n".join([
            f"[green]{i}. {hint['description']} ({hint['cost']})[/]" 
            for i, hint in enumerate(hints)
          ]))

          valid = False

          while not valid:
            hint_req = prompt(":light_bulb: [cyan]Please enter the hint requirements (comma separated)[/]")
            hint_req = [int(i) for i in hint_req.split(",")]
            # Check if the requirements are valid
            for requirement in hint_req:
              if requirement >= len(hints):
                console.print("[bright_red]Invalid hint requirement.[/bright_red]")
                break
            else:
              valid = True

      hints.append({
        "description": hint_description,
        "cost": cost,
        "requirements": hint_req
      })

      if not Confirm.ask(":light_bulb: [cyan]Do you want to add another hint?[/]"):
        break

      # Add spacing between the hints
      console.print()

    return hints
  else:
    return None
  

def get_files() -> list[Path] | None:
  """
  Prompts the user for the files.
  """
  console.rule(":file_folder: [bold yellow]Challenge Files[/] :file_folder:")

  has_files = Confirm.ask("[cyan]Are there any files that should be given to participants for this challenge?[/]")
  if has_files:
    return askopenfilenames(title="Please select the challenge files")
  else:
    return None
  

def get_requirements() -> list[str] | None:
  """
  Prompts the user for the requirements.
  """
  console.rule(":triangular_flag: [bold yellow]Challenge Requirements[/] :triangular_flag:")

  has_requirements = Confirm.ask("[cyan]Do any challenges need to be solved before this challenge?[/]")
  if has_requirements:
    requirements = []
    while True:
      requirement = prompt("[cyan]Please enter the challenge name of the requirement (case-insensitive)[/]")
      requirements.append(requirement)
      if not Confirm.ask("[cyan]Do you want to add another requirement?[/]"):
        break

      # Add spacing between the requirements
      console.print()
  else:
    requirements = None
  
  return requirements
  



@app.command()
def create_challenge_cli():
  """
  Creates a challenge folder in the current directory.
  """

  config = get_ctf_config()

  if config is None:
    return
  
  categories = config.categories
  difficulties = config.diff_names

  try:
    console.rule(":rocket: [bold yellow]Challenge Details[/] :rocket:")

    name = prompt(":rocket: [cyan][1/6] Please enter the challenge name (case-insensitive)[/]")

    description = get_description()
    if description is None:
      # User cancelled, abort
      console.print("[bright_red]Aborting...[/bright_red]")
      return

    # Print the categories in a list, with an index starting at 1
    console.print("\n[bright_yellow]Categories:[/]")
    for i, category in enumerate(categories):
      console.print(f"[bright_yellow]{i + 1}. {category.capitalize()}[/]")

    while True:
      category = prompt(f"\n:rocket: [cyan][3/6] Please choose the challenge category \[1-{len(categories)}][/]", method=IntPrompt.ask)
      if 1 > category > len(categories):
        console.print("[bright_red]Invalid category.[/bright_red]\n")
      else:
        category = categories[category - 1].capitalize()
        break

    console.print(f"[green]Category selected: {category}[/green]\n")


    # Print the difficulties in a list, with an index starting at 1
    console.print("[bright_yellow]Difficulties:[/]")
    for i, difficulty in enumerate(difficulties):
      console.print(f"[bright_yellow]{i + 1}. {difficulty.capitalize()}[/]")

    while True:
      difficulty = prompt(f"\n:rocket: [cyan][4/6] Please choose the challenge difficulty \[1-{len(difficulties)}][/]", method=IntPrompt.ask)
      if 1 > difficulty > len(difficulties):
        console.print("[bright_red]Invalid difficulty.[/bright_red]\n")
      else:
        difficulty = difficulties[difficulty - 1].capitalize()
        break

    console.print(f"[green]Difficulty selected: {difficulty}[/green]")

    author = prompt("\n:rocket: [cyan][5/6] Please enter your name (case-sensitive)[/]")

    discord = prompt("\n:rocket: [cyan][6/6] Please enter your discord username (So we can contact you if your challenge breaks)[/]")

    solution_files = get_solution_files()
    if solution_files == "":
      # User cancelled, abort
      console.print("[bright_red]No files selected, aborting...[/bright_red]")
      return
    
    elif solution_files is not None:
      solution_files = [Path(file) for file in solution_files]

      console.print("[green]Solution files selected: [/green]")
      for file in solution_files:
        console.print(f"[light_green]- {file}[/light_green]")


    flags = get_flags()

    hints = get_hints()

    requirements = get_requirements()

    files = get_files()
    if files == "":
      # User cancelled, abort
      console.print("[bright_red]No files selected, aborting...[/bright_red]")
      return
    
    elif files is not None:
      files = [Path(file) for file in files]

      console.print("[green]Files selected: [/green]")
      for file in files:
        console.print(f"[light_green]- {file}[/light_green]")


    console.rule(":file_folder: [bold yellow]Challenge Services[/] :file_folder:")
    has_serivices = Confirm.ask("[cyan]Does the challenge require hosting?[/]")
    if has_serivices:
      services = []

      # Add spacing
      console.print()

      while True:
        service_name = prompt("[cyan]Please enter the service name (this will be the name of the docker container)[/]")
        service_path = askdirectory(title="Please select the folder containing the dockerfile")

        if service_path == "":
          # User cancelled, abort
          console.print("[bright_red]No directory selected, aborting...[/bright_red]")
          return

        if not is_valid_service_folder(service_path):
          console.print("[bright_red]Invalid service folder, aborting...[/bright_red]")
          return

        service_port = prompt("[cyan]Please enter the port exposed by the service[/]", method=IntPrompt.ask)
        services.append({
          "name": service_name,
          "path": service_path,
          "port": service_port
        })
        if not Confirm.ask("[cyan]Do you want to add another service?[/]"):
          break
    else:
      services = None


    # console.print the challenge details in a panel
    info_string = f"[cyan][bold]Name:[/bold] {name}\n"
    info_string += f"[bold]Description:[/bold] {description}\n"
    info_string += f"[bold]Category:[/bold] {category}\n"
    info_string += f"[bold]Difficulty:[/bold] {difficulty}\n"
    info_string += f"[bold]Author:[/bold] {author}\n"
    info_string += f"[bold]Discord:[/bold] {discord}\n\n"

    info_string += f"[bold]Solution Files:[/bold]\n"
    for file in solution_files:
      info_string += f"- {file.as_posix()}\n"
    info_string += "\n"

    info_string += f"[bold]Flags:[/bold] {', '.join([flag['flag'] for flag in flags])}\n\n"

    info_string += f"[bold]Hints:[/bold] {hints}\n\n"

    info_string += f"[bold]Files:[/bold]\n"
    if files is not None:
      for file in files:
        info_string += f"- {file.as_posix()}\n"
    else:
      info_string += "None\n"
    info_string += "\n"

    info_string += f"[bold]Requirements:[/bold] {requirements}\n\n"
    info_string += f"[bold]Services:[/bold] {services}[/cyan]"


    console.print(Panel(
      info_string,
      border_style="green",
      title="[bright_yellow]Challenge Details[/bright_yellow]"
    ))


    # Confirm that the user wants to create the challenge
    if not Confirm.ask("[cyan]Do you want to create the challenge?[/]"):
      console.print("[bright_red]Aborting...[/bright_red]")
      return


    challenge_path = create_challenge(
      name=name,
      description=description,
      difficulty=difficulty,
      category=category,
      author=author,
      discord=discord,
      solution_files=solution_files,
      flags=flags,
      hints=hints,
      files=files,
      requirements=requirements,
      services=services
    )

    console.print(f"[green]Successfully created challenge `{name}` at `{challenge_path.resolve()}`[/green]")

  except (KeyboardInterrupt, EOFError):
    console.print("[bright_red]Aborting...[/bright_red]")
    return
  except Exception as e:
    # This should not happen
    console.print(f"[bright_red]An unexpected error occurred: {e}[/bright_red]")
    return
ctf_architect/chall_architect/cli.py
Jus-Codin/CTF-Architect
"""
Functions to create a new challenge folder.

Challenge template:
<challenge name>
├── dist
│   └── <files to give participants here>
├── service
│   └── <service name>
│       ├── dockerfile
│       └── <other files>
├── solution
│   ├── writeup.md
│   └── <optional solve scripts, etc>
├── chall.yaml
└── README.md
"""

from __future__ import annotations

import shutil
from pathlib import Path
from typing import TypedDict

from yaml import safe_dump

from ctf_architect.core.models import ChallengeInfo, Flag, Hint, Service


class ServiceDict(TypedDict):
  name: str
  path: str
  port: int


README_TEMPLATE = """# {name}
{description}

## Summary
- **Author:** {author}
- **Discord Username:** {discord}
- **Category:** {category}
- **Difficulty:** {difficulty}

## Hints
{hints}

## Files
{files}

## Services
{services}

## Flags
{flags}
"""


def create_challenge_readme(info: ChallengeInfo) -> str:
  """
  Creates a README.md file for a challenge.
  """
  flags = "\n".join([
    f"- `{flag.flag}` ({'regex' if flag.regex else 'static'})"
    for flag in info.flags
  ])
  
  if info.hints is None:
    hints = "None"
  else:
    hints = "\n".join([f"- `{hint.description}` ({hint.cost} points)" for hint in info.hints])
  
  if info.files is None:
    files = "None"
  else:
    files = ""
    for file in info.files:
      path = Path(file)
      files += f"- [`{path.name}`]({path})\n"
  
  if info.services is None:
    services = "None"
  else:
    services = ""
    for service in info.services:
      services += f"- [`{service.name}`]({service.path}) (port {service.port})\n"
  
  return README_TEMPLATE.format(
    name=info.name,
    description=info.description,
    author=info.author,
    discord=info.discord,
    category=info.category,
    difficulty=info.difficulty,
    hints=hints,
    files=files,
    services=services,
    flags=flags
  )


def create_challenge_folder(name: str, dist: bool = False, service: bool = False) -> Path:
  """
  Creates a folder with the challenge template.
  """
  challenge_path = Path(name)
  challenge_path.mkdir(exist_ok=True)

  if dist:
    (challenge_path / "dist").mkdir(exist_ok=True)

  if service:
    (challenge_path / "service").mkdir(exist_ok=True)

  (challenge_path / "solution").mkdir(exist_ok=True)

  # Create empty files
  (challenge_path / "chall.yaml").touch()
  (challenge_path / "README.md").touch()

  return challenge_path



def create_challenge(
  name: str,
  description: str,
  difficulty: str,
  category: str,
  author: str,
  discord: str,
  solution_files: list[Path],
  flag: str =  None,
  flags: list[dict[str, str | bool]] = None,
  hints: dict[str, str | int] = None,
  files: list[Path] = None,
  requirements: list[str] = None,
  services: list[ServiceDict] = None
) -> Path:
  """
  Creates a folder with the challenge template.

  Parameters
  ----------
  name : str
    The name of the challenge.
  description : str
    The description of the challenge.
  difficulty : str
    The difficulty of the challenge.
  category : str
    The category of the challenge.
  author : str
    The author of the challenge.
  discord : str
    The discord of the author.
  flag : str, optional
    The flag of the challenge.
  flags : list[dict[str, str | bool]], optional
    The flags of the challenge.
  hints : dict[str, str | int], optional
    The hints of the challenge.
  files : list[Path], optional
    Paths to the files that should be given to participants
  requirements : list[str], optional
    The requirements of the challenge.
  services : list[ServiceDict], optional
    The services of the challenge.

  Returns
  -------
  Path
    The path to the challenge folder.
  """

  if flag is None and flags is None:
    raise ValueError("Must specify either flag or flags")
  elif flag is not None and flags is not None:
    raise ValueError("Cannot specify both flag and flags")
  
  if flag is not None:
    flags = [Flag(flag=flag)]

  if flags is not None:
    flags = [Flag(**flag) for flag in flags]

  if hints is not None:
    hints = [Hint(**hint) for hint in hints]
  
  dist = files is not None
  service = services is not None

  path = create_challenge_folder(name, dist=dist, service=service)

  if files is not None:
    # Copy files to dist folder
    file_paths = []
    for file in files:
      file_path = shutil.copy(file, path / "dist")
      file_path = Path(file_path).relative_to(path)
      file_paths.append(file_path.as_posix())
  else:
    file_paths = None

  if services is not None:
    # Copy services to service folder
    service_objs = []
    for service in services:
      service_path = path / "service" / service["name"]
      shutil.copytree(service["path"], service_path)
      service_obj = Service(service["name"], path=service_path.relative_to(path).as_posix(), port=str(service["port"]))
      service_objs.append(service_obj)
  else:
    service_objs = None

  # Copy solution files to solution folder
  for file in solution_files:
    shutil.copy(file, path / "solution")

  # Create challenge info
  info = ChallengeInfo(
    name=name,
    description=description,
    difficulty=difficulty,
    category=category,
    author=author,
    discord=discord,
    flags=flags,
    hints=hints,
    files=file_paths,
    requirements=requirements,
    services=service_objs
  )

  # Save challenge info in chall.yaml
  safe_dump(info.to_dict(), (path / "chall.yaml").open("w"))

  # Create README.md
  (path / "README.md").write_text(create_challenge_readme(info))

  return path
ctf_architect/chall_architect/create.py
Jus-Codin/CTF-Architect
"""
Various utility functions for chall-architect
"""

from __future__ import annotations

from pathlib import Path

from yaml import safe_load

from ctf_architect.core.models import Config


def get_config(path: str | Path) -> Config:
  """
  Gets the config from a given path.
  """
  if isinstance(path, str):
    path = Path(path)

  # Ensure path is a file
  if not path.is_file():
    raise ValueError("Path must be a file")
  
  # Ensure path is a yaml file
  if path.suffix != ".yaml":
    raise ValueError("Path must be a yaml file")
  
  # Load the config
  with path.open("r") as f:
    data = safe_load(path.open("r"))

  config = data.get("config")

  # Ensure config is present
  if config is None:
    raise ValueError("Config not found in yaml file")
  
  config = Config(**config)

  return config


def is_valid_service_folder(path: str | Path) -> bool:
  """
  Checks if the given path is a valid service folder.

  A service folder is considered valid if it has a dockerfile.
  """
  if isinstance(path, str):
    path = Path(path)

  # Check if there is a docker file in the folder, case-insensitive
  return any(file.name.lower() == "dockerfile" for file in path.iterdir())
ctf_architect/chall_architect/utils.py
Jus-Codin/CTF-Architect
from __future__ import annotations

from pathlib import Path

import typer
from rich.console import Console
from rich.panel import Panel
from rich.prompt import Confirm, IntPrompt, Prompt

from ctf_architect.cli.challenge import challenge_app
from ctf_architect.cli.mapping import mapping_app
from ctf_architect.cli.stats import stats_app
from ctf_architect.core.initialize import init_no_config, init_with_config

console = Console()


app = typer.Typer()
app.add_typer(challenge_app, name="challenge")
app.add_typer(mapping_app, name="mapping")
app.add_typer(stats_app, name="stats")


@app.command()
def init(
  port: int = typer.Option(8000, "--port", "-p", help="Specify the port."),
  config_only: bool = typer.Option(False, "--config-only", "-c", help="Only create the config file.")
):
  """
  Initialize a new CTF repo.
  """

  if config_only and Path("ctf_config.yaml").exists():
    console.print("[bright_red]ctf_config.yaml already exists.[/bright_red]")
    return
  
  # Check if there is a ctf_config.yaml in the current directory, if so, use that config
  # Otherwise, use the config specified in the command line arguments
  if Path("ctf_config.yaml").exists():
    try:
      init_with_config()
    except ValueError as e:
      console.print(f"[bright_red]{e}[/bright_red]")
      return
    else:
      console.print("[bright_green]Initialized CTF repo with the config in ctf_config.yaml.[/bright_green]")
  
  else:

    # Get the categories and difficulties from the user
    categories = []
    difficulties = []

    try:
      console.rule("[bold yellow]CTF Categories[/]")

      console.print("[bold cyan]Please enter the categories for your CTF. Enter an empty string to stop.[/]")

      while True:
        category = Prompt.ask("[cyan]Category name[/]", default="", show_default=False)
        if category == "":
          break
        categories.append(category.lower())

        console.print(f"\n[green]Category {category} added.[/]\n")

      if len(categories) == 0:
        console.print("[bright_red]Please specify at least one category.[/bright_red]")
        return
      else:
        console.print(f"\n[bold green]Categories:[/]")
        for category in categories:
          console.print(f"[green] - {category}[/]")


      console.rule("[bold yellow]CTF Difficulties[/]")

      console.print("\n[bold cyan]Please enter the difficulties for your CTF. Enter an empty string to stop.[/]")

      while True:
        name = Prompt.ask("[cyan]Difficulty name[/]", default="", show_default=False)
        if name == "":
          break
        value = IntPrompt.ask("[cyan]Difficulty value[/]", default=500)
        difficulties.append({
          "name": name.lower(),
          "value": value
        })

        console.print(f"\n[green]Difficulty {name} added. ({value} points)[/]\n")

      if len(difficulties) == 0:
        console.print("[bright_red]Please specify at least one difficulty.[/bright_red]")
        return
      else:
        console.print(f"\n[bold green]Difficulties:[/]")
        for difficulty in difficulties:
          console.print(f"[green] - {difficulty['name']} ({difficulty['value']} points)[/]")
    except (KeyboardInterrupt, EOFError):
      console.print("\n[bright_red]Aborting...[/bright_red]")
      return


    # Show panel of categories and difficulties
    # If the user confirms, initialize the CTF repo
    
    config_text = ""
    config_text += "[bold cyan]Categories:[/]\n"
    for category in categories:
      config_text += f"[cyan] - {category}[/]\n"
    config_text += "\n"
    config_text += "[bold cyan]Difficulties:[/]\n"
    for difficulty in difficulties:
      config_text += f"[cyan] - {difficulty['name']} ({difficulty['value']} points)[/]\n"

    console.print(
      Panel(
        config_text,
        title="[bright_yellow]CTF Config[/bright_yellow]",
        border_style="green"
      )
    )

    if not Confirm.ask("[cyan]Is this correct?[/]"):
      console.print("[bright_red]Aborting...[/bright_red]")
      return

    
    try:
      init_no_config(categories, difficulties, port, config_only)
    except ValueError as e:
      console.print(f"[bright_red]{e}[/bright_red]")
      return
    else:
      console.print("[bright_green]Initialized CTF repo with specified config.[/bright_green]")
ctf_architect/cli/__init__.py
Jus-Codin/CTF-Architect
from __future__ import annotations

from pathlib import Path
import shutil

import typer
from rich import print
from rich.panel import Panel

from ctf_architect.core.challenge import add_challenge, find_challenge
from ctf_architect.core.config import load_config
from ctf_architect.core.constants import APP_CMD_NAME
from ctf_architect.core.stats import (update_category_readme,
                         update_root_readme)
from ctf_architect.core.unzip import unzip


def is_challenge_repo() -> bool:
  """
  Checks if the current working directory is a challenge repo.

  A challenge repo is considered valid if it has a ctf_config.yaml and a challenges folder.
  """
  return Path("ctf_config.yaml").exists() and Path("challenges").exists()


challenge_app = typer.Typer()


@challenge_app.callback()
def callback():
  """
  Commands to manage challenges.
  """


@challenge_app.command("import")
def challenge_import(
  path: str = typer.Option(default="./", help="Specify the path to the challenge zip files."),
  delete_on_error: bool = typer.Option(False, "--delete-on-error", "-d", help="Delete the challenge folder if an error occurs."),
  replace: bool = typer.Option(False, "--replace", "-r", help="Replace all existing challenges."),
  update_stats: bool = typer.Option(True, "--update-stats", "-u", help="Update the stats after importing.")
):
  """
  Imports challenges from zip files in the specified path.
  """
  try:
    config = load_config()
  except FileNotFoundError:
    print(f"[bright_red]Challenge repo not found, are you in the correct directory? If so, please run `{APP_CMD_NAME} init` first.")
    return
  
  zip_path = Path(path)
  extracted_path = Path("extracted")

  if not zip_path.exists():
    print(f"[bright_red]Path `{zip_path}` does not exist.[/bright_red]")
    return
  
  success = 0
  unzip_failed = 0
  add_failed = 0

  for zip_file in zip_path.glob("*.zip"):
    try:
      # Unzip the file
      extracted = unzip(zip_file, extracted_path, delete=False)
    except Exception as e:
      print(f"[bright_red]Failed to unzip challenge `{zip_file.name}`:[/bright_red]")
      print(f"[bright_red]{e}[/bright_red]")
      unzip_failed += 1
      continue
    
    try:
      # Add the challenge
      add_challenge(extracted, replace=replace)
    except FileExistsError:
      # Prompt the user if they want to replace the challenge
      print(f"[bright_red]Challenge `{extracted.name}` already exists.[/bright_red]")
      if typer.confirm("Do you want to replace it?"):
        add_challenge(extracted, replace=True)
        print(f"[green]Successfully added challenge `{extracted.name}`[/green]")
        success += 1
      else:
        print(f"[bright_red]Skipping challenge `{extracted.name}`[/bright_red]")
        add_failed += 1
        continue
    except Exception as e:
      print(f"[bright_red]Failed to add challenge `{extracted.name}`:[/bright_red]")
      print(f"[bright_red]{e}[/bright_red]")
      # Delete the challenge folder if it exists
      if extracted.exists() and delete_on_error:
        shutil.rmtree(extracted)
      add_failed += 1
      continue
    else:
      print(f"[green]Successfully added challenge `{extracted.name}`[/green]")
      success += 1

  print(f"[green]Successfully imported {success} challenges[/green]")
  if unzip_failed > 0:
    print(f"[bright_red]Failed to unzip {unzip_failed} challenges[/bright_red]")
  if add_failed > 0:
    print(f"[bright_red]Failed to add {add_failed} challenges[/bright_red]")

  # Update the stats
  if update_stats and success > 0:
    for category in config.categories:
      update_category_readme(category)
    
    update_root_readme()

    print("[green]Repo READMEs updated successfully.")


@challenge_app.command("export")
def challenge_export(
  name: str = typer.Argument(..., help="The name of the challenge to export."),
  path: str = typer.Option("./", "--path", "-p", help="Specify the path to export the challenge to."),
  filename: str = typer.Option(None, "--filename", "-f", help="Specify the filename of the exported challenge."),
):
  """
  Exports the challenge as a zip file

  NOTE: For some reason if a directory is created where the zip file is supposed to be, admin privileges are required to delete the directory.
  """
  if not is_challenge_repo():
    print(f"[bright_red]Challenge repo not found, are you in the correct directory? If so, please run `{APP_CMD_NAME} init` first.")
    return
  
  challenge = find_challenge(name)

  if challenge is None:
    print(f"[bright_red]Challenge `{name}` not found.[/bright_red]")
    return
  
  challenge_path = challenge[1]

  if filename is None:
    filename = challenge_path.name

  path = Path(path) / filename

  archive_path = shutil.make_archive(path, "zip", challenge_path)

  print(f"[green]Successfully exported challenge `{name}` to `{archive_path}`[/green]")


@challenge_app.command("info")
def challenge_info(name: str = typer.Argument(..., help="The name of the challenge to get info for.")):
  """
  Gets the info of a challenge.
  """
  try:
    config = load_config()
  except FileNotFoundError:
    print(f"[bright_red]Challenge repo not found, are you in the correct directory? If so, please run `{APP_CMD_NAME} init` first.")
    return
  
  challenge = find_challenge(name)

  if challenge is None:
    print(f"[bright_red]Challenge `{name}` not found.[/bright_red]")
    return
  
  challenge_info, challenge_path = challenge

  info = (
    f"[bold]Name:[/] {challenge_info.name}\n"
    f"[bold]Description:[/] {challenge_info.description.strip()}\n\n"

    f"[bold]Difficulty:[/] {challenge_info.difficulty.capitalize()}\n"
    f"[bold]Category:[/] {challenge_info.category.capitalize()}\n"
    f"[bold]Author:[/] {challenge_info.author}\n"
    f"[bold]Discord:[/] {challenge_info.discord}\n"
  )

  if challenge_info.services is not None:
    info += "\n[bold]Services:[/]"
    for service in challenge_info.services:
      info += (
        f"\n - [bold]Name:[/] {service.name}\n"
        f" - [bold]Path:[/] {service.path}\n"
        # f" - [bold]Path:[/] {(challenge_path / service.path).resolve()}\n"
        f" - [bold]Port:[/] {service.port}\n"
      )

  print(Panel.fit(
    info,
    title=f"Challenge Info for {challenge_info.name}",
    border_style="bright_cyan",
    padding=(1, 2)
  ))
ctf_architect/cli/challenge.py
Jus-Codin/CTF-Architect
from __future__ import annotations

import typer
from rich import print
from rich.panel import Panel
from rich.prompt import Confirm
from rich.table import Table

from ctf_architect.core.config import load_config, save_config
from ctf_architect.core.constants import APP_CMD_NAME
from ctf_architect.core.mappings import generate_port_mappings

mapping_app = typer.Typer()


def create_mapping_table(mappings: dict) -> Table:
  """
  Creates a rich table for the port mappings.
  """
  table = Table(title="Port Mappings")
  table.add_column("Service Name", header_style="bright_cyan", style="cyan")
  table.add_column("Container Port", header_style="bright_green", style="green")
  table.add_column("Host Port", header_style="bright_green", style="green")

  for service_name, port_mapping in mappings.items():
    table.add_row(service_name, port_mapping["from_port"], port_mapping["to_port"])

  return table


@mapping_app.callback()
def callback():
  """
  Commands to manage port mappings for challenges
  """


@mapping_app.command("generate")
def generate():
  """
  Generates the port mappings for challenges.
  """
  try:
    config = load_config()
  except FileNotFoundError:
    print(f"[bright_red]Challenge repo not found, are you in the correct directory? If so, please run `{APP_CMD_NAME} init` first.")
    return
  
  mapping = generate_port_mappings()

  # Print and display the mappings in a table
  table = create_mapping_table(mapping)

  print(table)

  # Ask for confirmation to save the mappings
  if not Confirm.ask("Save port mappings to ctf_config.yaml?"):
    return

  config.port_mappings = mapping
  save_config(config)

  print("[green]Port mappings saved successfully.")


@mapping_app.command("show")
def show():
  # Show port mappings
  try:
    config = load_config()
  except FileNotFoundError:
    print(f"[bright_red]Challenge repo not found, are you in the correct directory? If so, please run `{APP_CMD_NAME} init` first.")
    return
  
  if config.port_mappings is None:
    print("[bright_red]No port mappings found.[/bright_red]")
    return
  
  table = create_mapping_table(config.port_mappings)

  print(table)
ctf_architect/cli/mapping.py
Jus-Codin/CTF-Architect
from __future__ import annotations

import typer
from rich import print
from rich.table import Table

from ctf_architect.core.config import load_config
from ctf_architect.core.constants import APP_CMD_NAME
from ctf_architect.core.stats import (get_category_diff_stats, update_category_readme,
                         update_root_readme)


stats_app = typer.Typer()


@stats_app.callback()
def callback():
  """
  Commands to manage the stats of the challenge repo.
  """


@stats_app.command("update")
def stats_update():
  """
  Updates all the stats in the READMEs.
  """
  try:
    config = load_config()
  except FileNotFoundError:
    print(f"[bright_red]Challenge repo not found, are you in the correct directory? If so, please run `{APP_CMD_NAME} init` first.")
    return

  try:
    for category in config.categories:
      update_category_readme(category)
    
    update_root_readme()
  except FileNotFoundError:
    print(f"[bright_red]README not found. Please run `{APP_CMD_NAME} init` first.")
  else:
    print("[green]READMEs updated successfully.")


@stats_app.command("show")
def stats_show(category: str = typer.Argument(None)):
  """
  Shows the stats of a category.

  If no category is specified, shows the stats of all categories.
  """
  try:
    config = load_config()
  except FileNotFoundError:
    print(f"[bright_red]Challenge repo not found, are you in the correct directory? If so, please run `{APP_CMD_NAME} init` first.")
    return

  if category is None:
    table = Table(title="Overall Difficulty Distribution")
    # Create a table in the following format
    # | Category | Easy | Medium | ... | Total |
    # | -------- | ---- | ------ | --- | ----- |
    # | ...      | ...  | ...    | ... | ...   |
    # | ...      | ...  | ...    | ... | ...   |
    # | Total    | ...  | ...    | ... | ...   |
    table.add_column("Category", header_style="bright_cyan", style="cyan", no_wrap=True)
    for difficulty in config.diff_names:
      table.add_column(difficulty.capitalize(), header_style="magenta", style="magenta", justify="center")
    table.add_column("Total", header_style="bright_yellow", style="yellow", justify="center")

    stats = {category: get_category_diff_stats(category) for category in config.categories}

    for category in stats:
      # Check if it is the last category to add to the list
      is_last = category == config.categories[-1]
      table.add_row(
        category.capitalize(),
        *[str(stats[category][difficulty]) for difficulty in config.diff_names],
        str(sum(stats[category].values())),
        end_section=is_last
      )

    # Add the total row
    total_row = []
    for difficulty in config.diff_names:
      total_row.append(sum(stats[category][difficulty] for category in config.categories))

    table.add_row(
      "Total",
      *map(str, total_row),
      str(sum(total_row))
    )
    
  else:
    table = Table(title=f"{category.capitalize()} Difficulty Distribution")
    # Create a table in the following format
    # | Difficulty | Count |
    # | ---------- | ----- |
    # | Easy       | ...   |
    # | Medium     | ...   |
    # | ...        | ...   |
    table.add_column("Difficulty", header_style="bright_cyan", style="cyan")
    table.add_column("Count", header_style="bright_yellow", style="yellow", justify="center")

    stats = get_category_diff_stats(category)

    for difficulty, count in stats.items():
      table.add_row(difficulty.capitalize(), str(count))

  print(table)
ctf_architect/cli/stats.py
Jus-Codin/CTF-Architect
from __future__ import annotations

import os
import shutil
from pathlib import Path
from yaml import safe_load

from typing import Generator

from ctf_architect.core.config import load_config
from ctf_architect.core.models import ChallengeInfo, Service


def is_valid_challenge(path: str | Path) -> bool:
  """
  Checks if a challenge is valid.
  A challenge is considered valid if it has a chall.yaml and a README.md file.
  """
  files = os.listdir(path)
  if "chall.yaml" not in files:
    return False
  if "README.md" not in files:
    return False
  
  return True


def get_challenge_info(path: str | Path) -> ChallengeInfo:
  if isinstance(path, str):
    path = Path(path)

  if not is_valid_challenge(path):
    raise ValueError("Challenge does not have chall.yaml or README.md")
  
  with (path / "chall.yaml").open("r") as f:
    data = safe_load(f)

  info = data.get("challenge")
  if info is None:
    raise ValueError("Challenge does not specify challenge info")
  
  services = data.get("services")

  if services is not None:
    services = [
      Service(**v)
      for _, v in services.items()
    ]
  
  return ChallengeInfo.from_dict(info, services=services)


def get_services(path: str | Path) -> list[Service] | None:
  if isinstance(path, str):
    path = Path(path)

  if not is_valid_challenge(path):
    raise ValueError("Invalid challenge")
  
  with (path / "chall.yaml").open("r") as f:
    data = safe_load(f)
  
  services = data.get("services")
  if services is None:
    return None
  else:
    return [
      Service(name=k, **v)
      for k, v in services.items()
    ]


def verify_challenge_info(challenge: ChallengeInfo):
  """
  Verifies that the challenge info has valid categories and difficulties.
  """

  config = load_config()

  if challenge.category.lower() not in config.categories:
    raise ValueError(f"Invalid category: {challenge.category}")
  
  if challenge.difficulty.lower() not in config.diff_names:
    raise ValueError(f"Invalid difficulty: {challenge.difficulty}")
  

def find_challenge(name: str) -> tuple[ChallengeInfo, Path] | None:
  """
  Tries to find a challenge with the given name.

  Search method:
  1. Search by the folder name, if substring match found, check chall.yaml to verify
  2. Search every chall.yaml for a name match

  Returns a tuple of the challenge info and the path to the challenge folder.
  """
  config = load_config()

  challenges_path = Path("./challenges")

  folders: list[Path] = []
  for category in config.categories:
    for dir in (challenges_path / category).iterdir():
      if dir.is_dir():
        folders.append(dir)

  folders_searched = []

  for folder in folders:
    if name.lower() in folder.name.lower():
      folders_searched.append(folder)
      try:
        challenge = get_challenge_info(folder)
        return challenge, folder
      except ValueError:
        pass

  for folder in folders:
    if folder in folders_searched:
      continue
    try:
      challenge = get_challenge_info(folder)
      if name.lower() == challenge.name.lower():
        return challenge, folder
    except ValueError:
      pass


def add_challenge(folder: str | Path, replace: bool = False):
  """
  Adds a challenge to the challenges directory.

  If replace is True, the challenge will be replaced if it already exists.
  """
  if isinstance(folder, str):
    folder = Path(folder)

  if not folder.is_dir():
    raise ValueError("Specified path is not a directory!")
  
  challenge = get_challenge_info(folder)
  verify_challenge_info(challenge)

  # Check if challenge already exists
  if (c := find_challenge(challenge.name)) is not None:
    if replace:
      remove_challenge(path=c[1])
    else:
      raise FileExistsError("Challenge already exists!")
  
  
  # Move the challenge to the correct category
  new_path = Path("./challenges") / challenge.category.lower() / challenge.name
  folder.rename(new_path)


def remove_challenge(name: str | None = None, path: Path | str | None = None):
  """
  Removes a challenge

  Either challenge name or path must be specified.

  If challenge name is specified, the challenge will be searched for and removed.
  """
  if name is None and path is None:
    raise ValueError("Either challenge name or path must be specified!")
  
  if name is not None:
    challenge = find_challenge(name)
    if challenge is None:
      raise ValueError("Challenge not found!")
    path = challenge[1]

  if isinstance(path, str):
    path = Path(path)

  if not path.is_dir():
    raise ValueError("Specified path is not a directory!")
  
  shutil.rmtree(path)


def walk_challenges(category: str | None = None) -> Generator[ChallengeInfo]:
  """
  Walks through all challenges in the challenges directory.

  Can optionally specify a category to only walk through challenges in that category.
  """
  config = load_config()

  challenges_path = Path("./challenges")

  for _category in config.categories:
    if category is not None and category.lower() != _category.lower():
      continue
    for dir in (challenges_path / _category).iterdir():
      if dir.is_dir():
        try:
          yield get_challenge_info(dir)
        except ValueError:
          pass
ctf_architect/core/challenge.py
Jus-Codin/CTF-Architect
import os
from yaml import safe_load, safe_dump

from ctf_architect.core.models import Config


def load_config() -> Config:
  files = os.listdir("./")
  if "ctf_config.yaml" not in files:
    raise FileNotFoundError("ctf_config.yaml not found")
  
  with open("ctf_config.yaml", "r") as f:
    data = safe_load(f)

  config = data.get("config")
  if config is None:
    raise ValueError("ctf_config.yaml does not specify config")

  return Config(**config)


def save_config(config: Config):
  with open("ctf_config.yaml", "w") as f:
    data = {
      "config": config.to_dict()
    }
    
    safe_dump(data, f)
ctf_architect/core/config.py
Jus-Codin/CTF-Architect
APP_CMD_NAME = "ctf-architect"
ctf_architect/core/constants.py
Jus-Codin/CTF-Architect
"""
Functionality for creating a docker compose file
"""

from __future__ import annotations

from pathlib import Path
from yaml import safe_dump

from ctf_architect.core.challenge import walk_challenges
from ctf_architect.core.config import load_config
from ctf_architect.core.models import Service


def create_compose_dict(
    service: Service,
    challenge_name: str,
    category: str,
    host_port: str
  ) -> dict:
  """
  Creates a docker compose service for a challenge service
  """
  
  service_path = (
    Path("./challenges") / category /
    challenge_name / service.path
  ).as_posix()

  d = {
      "build": service_path,
      "ports": [f"{host_port}:{service.port}"],
      "container_name": service.name,
  }

  if service.extras is not None:
    d.update(service.extras)

  if "restart" not in d:
    d["restart"] = "always"

  return d


def create_compose_file() -> None:
  """
  Creates a docker compose file for all challenges
  """

  config = load_config()

  if config.port_mappings is None:
    raise ValueError("Port mappings not found in config, please generate them first")

  services = {}

  for challenge in walk_challenges():
    if challenge.services is not None:
      for service in challenge.services:
        services[service.name] = create_compose_dict(
          service,
          challenge.name,
          challenge.category.lower(),
          config.port_mappings[service.name]["to_port"]
        )

  with open("docker-compose.yml", "w") as f:
    safe_dump({
      "version": "3",
      "services": services
    }, f)
ctf_architect/core/docker.py
Jus-Codin/CTF-Architect
"""
Contains the code to initialize a new CTF repo
"""

from __future__ import annotations

from pathlib import Path

from ctf_architect.core.config import load_config, save_config
from ctf_architect.core.models import Config
from ctf_architect.core.stats import update_category_readme, update_root_readme


def init_no_config(
  categories: list[str],
  difficulties: list[dict[str, str | int]],
  port: int,
  config_only: bool = False
):
  """
  Initialize the repo with the given config.
  """
  config = Config(
    categories=categories,
    difficulties=difficulties,
    port=port
  )

  save_config(config)

  if not config_only:
    # Create the challenge folder
    Path("challenges").mkdir(exist_ok=True)

    # Create the folders for each category
    for category in categories:
      (Path("challenges") / category.lower()).mkdir(exist_ok=True)

    update_root_readme()

    for category in categories:
      update_category_readme(category)


def init_with_config():
  """
  Initialize the repo with the config in ctf_config.yaml
  """
  config = load_config()

  # Create the challenge folder
  Path("challenges").mkdir(exist_ok=True)

  # Create the folders for each category
  for category in config.categories:
    (Path("challenges") / category.lower()).mkdir(exist_ok=True)

  update_root_readme()

  for category in config.categories:
    update_category_readme(category)
ctf_architect/core/initialize.py
Jus-Codin/CTF-Architect
from __future__ import annotations

from ctf_architect.core.config import load_config
from ctf_architect.core.challenge import walk_challenges
from ctf_architect.core.models import PortMappingDict


def generate_port_mappings() -> dict[str, PortMappingDict]:
  config = load_config()

  port = config.port

  mapping: dict[str, PortMappingDict] = {}

  for category in config.categories:
    for challenge in walk_challenges(category):
      if challenge.services is not None:
        for service in challenge.services:
          if service.name in mapping:
            raise ValueError(f"Duplicate service name: {service.name}")
          mapping[service.name] = {
            "from_port": service.port,
            "to_port": str(port)
          }
          port += 1
    # If there is at least 1 service, go to the next thousand
    if port % 1000:
      port += 1000 - (port % 1000)

  # Check if we exceeded the port range
  # I'm sure this won't happen but just in case
  if port > 65535:
    raise ValueError("Port range exceeded")

  return mapping
ctf_architect/core/mappings.py
Jus-Codin/CTF-Architect
from __future__ import annotations

from dataclasses import dataclass
from typing import TypedDict


@dataclass
class Flag:
  flag: str
  regex: bool = False

  def to_dict(self) -> dict:
    return {
      "flag": self.flag,
      "regex": self.regex
    }


@dataclass
class Hint:
  description: str
  cost: int
  requirements: list[int] | None = None

  def to_dict(self) -> dict:
    return {
      "description": self.description,
      "cost": self.cost,
      "requirements": self.requirements
    }


@dataclass
class Service:
  name: str
  path: str
  port: str
  extras: dict[str, str] | None = None

  def to_dict(self) -> dict:
    d = {
      "name": self.name,
      "path": self.path,
      "port": self.port
    }
    if self.extras is not None:
      d["extras"] = self.extras
    return d


@dataclass
class ChallengeInfo:
  name: str
  description: str
  difficulty: str
  category: str
  author: str
  discord: str
  flags: list[Flag]

  hints: list[Hint] | None = None
  files: list[str] | None = None
  requirements: list[str] | None = None
  
  services: list[Service] | None = None

  @classmethod
  def from_dict(cls, d: dict, services: list[Service] | None = None) -> ChallengeInfo:
    flags = []
    for flag in d.pop("flags"):
      if isinstance(flag, str):
        flags.append(Flag(flag))
      else:
        flags.append(Flag(**flag))

    hints = d.pop("hints")
    if hints is not None:
      hints = [Hint(**hint) for hint in hints]

    return cls(
      flags=flags,
      hints=hints,
      services=services,
      **d
    )
  
  def to_dict(self) -> dict[str, dict | list[dict]]:
    return {
      "challenge": {
        "name": self.name,
        "description": self.description,
        "difficulty": self.difficulty,
        "category": self.category,
        "author": self.author,
        "discord": self.discord,
        "flags": [flag.to_dict() for flag in self.flags],
        "hints": [hint.to_dict() for hint in self.hints] if self.hints is not None else None,
        "files": self.files,
        "requirements": self.requirements
      },
      "services": {
        service.name: service.to_dict()
        for service in self.services
      } if self.services is not None else None
    }


class PortMappingDict(TypedDict):
  from_port: str
  to_port: str


@dataclass
class Config:
  categories: list[str]
  difficulties: list[dict[str, str | int]]
  port: int
  name: str | None = None
  port_mappings: dict[str, PortMappingDict] | None = None

  @property
  def diff_names(self) -> list[str]:
    return [d["name"] for d in self.difficulties]
  
  def to_dict(self) -> dict:
    return {
      "categories": self.categories,
      "difficulties": self.difficulties,
      "port": self.port,
      "name": self.name,
      "port_mappings": self.port_mappings
    }
ctf_architect/core/models.py
Jus-Codin/CTF-Architect
from __future__ import annotations

from pathlib import Path

from ctf_architect.core.challenge import get_challenge_info
from ctf_architect.core.config import load_config


CATEGORY_README_TEMPLATE = """# {name} Challenges
This directory contains challenges related to {name}.

## Difficulty Distribution
| Difficulty | Number of Challenges |
| ---------- |:--------------------:|
{diff_table}

## Challenges ({count} total)
| Name | Description | Difficulty | Author |
| ---- | ----------- | ---------- | ------ |
{challenges_table}
"""

ROOT_README_TEMPLATE = """# Challenges
This directory contains all challenges.

## Difficulty Distribution
{diff_table}

## Challenges ({count} total)
| Name | Description | Category | Difficulty | Author |
| ---- | ----------- | -------- | ---------- | ------ |
{challenges_table}
"""


def get_category_diff_stats(name: str) -> dict[str, int]:
  """
  Gets the difficulty distribution of a category.
  """
  config = load_config()

  if name.lower() not in config.categories:
    raise ValueError(f"Category {name} does not exist.")
  
  category_path = Path("challenges") / name.lower()

  difficulties = config.difficulties

  stats = {difficulty["name"]: 0 for difficulty in difficulties}

  for challenge_path in category_path.iterdir():
    if challenge_path.is_dir():
      info = get_challenge_info(challenge_path)
      stats[info.difficulty.lower()] += 1

  return stats


def update_category_readme(name: str):
  """
  Updates the category's README.md file.
  """
  config = load_config()

  if name.lower() not in config.categories:
    raise ValueError(f"Category {name} does not exist.")
  
  challenges_path = Path("challenges")

  # Check if the challenges directory exists
  if not challenges_path.exists():
    raise FileNotFoundError("Challenges directory does not exist, are you sure you are in the right directory?")
  
  category_path = challenges_path / name.lower()

  # Get all challenge info
  challenges: list[list[str]] = []
  
  for challenge_path in category_path.iterdir():
    if challenge_path.is_dir():
      info = get_challenge_info(challenge_path)
      challenges.append([info.name, info.description, info.difficulty, info.author])

  # Get the difficulty stats
  stats = get_category_diff_stats(name)

  # Create the challenges table
  challenges_table = "\n".join(
    f"| [{name}](<./{name}>) |" +
    f" {description if len(description) <= 20 else description[:20]+'...'} |" +
    f" {difficulty.capitalize()} |" +
    f" {author} |"
    for name, description, difficulty, author in challenges
  )

  # Create the difficulty table
  diff_table = "\n".join(
    f"| {difficulty.capitalize()} | {count} |"
    for difficulty, count in stats.items()
  )

  # Add the total row
  diff_table += f"\n| Total | {sum(stats.values())} |"

  # Create the README
  readme = CATEGORY_README_TEMPLATE.format(
    name=name.capitalize(),
    count=len(challenges),
    challenges_table=challenges_table,  
    diff_table=diff_table
  )

  # Write the README
  with (category_path / "README.md").open("w") as f:
    f.write(readme)


def update_root_readme():
  """
  Updates the root challenges README.md file.
  """
  config = load_config()

  challenges_path = Path("challenges")

  # Check if the challenges directory exists
  if not challenges_path.exists():
    raise FileNotFoundError("Challenges directory does not exist, are you sure you are in the right directory?")

  # Get all challenge info
  challenges: list[list[str]] = []
  
  for category in config.categories:
    category_path = challenges_path / category.lower()

    for challenge_path in category_path.iterdir():
      if challenge_path.is_dir():
        info = get_challenge_info(challenge_path)
        challenges.append([info.name, info.description, info.category, info.difficulty, info.author])

  # Get the difficulty stats
  stats = {category: get_category_diff_stats(category) for category in config.categories}

  # Create the challenges table
  challenges_table = "\n".join(
    f"| [{name}](<./{category.lower()}/{name}>) |" +
    f" {description if len(description) <= 20 else description[:20]+'...'} |" +
    f" {category.capitalize()} |" +
    f" {difficulty.capitalize()} |" +
    f" {author} |"
    for name, description, category, difficulty, author in challenges
  )

  # Create the difficulty table
  diff_table_header = "| Category | " + " | ".join(d.capitalize() for d in config.diff_names) + " | Total |\n"
  diff_table_header += "| -------- |:" + ":|:".join("-" * len(diff) for diff in config.diff_names) + ":|:-----:|\n"

  diff_table_body = "\n".join(
    f"| {category.capitalize()} | " + " | ".join(str(stats[category][diff]) for diff in config.diff_names) + f" | {sum(stats[category].values())} |"
    for category in stats
  )

  diff_table_total = (
    "\n| Total | " +
    " | ".join(str(sum(stats[category][diff] for category in config.categories)) for diff in config.diff_names) +
    f" | {sum(sum(stats[category].values()) for category in config.categories)} |\n"
  )

  diff_table = diff_table_header + diff_table_body + diff_table_total

  # Create the README
  readme = ROOT_README_TEMPLATE.format(
    count=len(challenges),
    challenges_table=challenges_table,  
    diff_table=diff_table
  )

  # Write the README
  with (challenges_path / "README.md").open("w") as f:
    f.write(readme)
ctf_architect/core/stats.py
Jus-Codin/CTF-Architect
from pathlib import Path
import zipfile


def unzip(file: Path, to: Path, delete: bool = True) -> Path:
  """
  Unzips a zip file to the specified path.
  """
  new_path = to / file.name[:-4]

  with zipfile.ZipFile(file, "r") as zip_ref:
    zip_ref.extractall(new_path)

  # move subfolders to the root
  while len(list(new_path.iterdir())) == 1:
    subfolder = list(new_path.iterdir())[0]
    if not subfolder.is_dir():
      break
    for f in subfolder.iterdir():
      f.rename(new_path / f.name)

    subfolder.rmdir()
    
  if delete:
    # Delete the zip file
    file.unlink()

  return new_path


def unzip_all(fp: str = "./", to: str = "./extracted/", delete: bool = True) -> list[Path]:
  """
  Unzips all zip files in the specified path.
  """
  path = Path(fp)
  to = Path(to)
  new_paths = []
  for file in path.glob("*.zip"):
    try:
      new_paths.append(unzip(file, to, delete=delete))
    except Exception as e:
      print(f"Error unzipping {file}: {e}")

  return new_paths
ctf_architect/core/unzip.py
Jus-Codin/CTF-Architect