Merge pull request #24 from dreadnode/simone/eng-458-cli-move-to-manifests-for-agent-templates

eng-458: implemented agent templates system

Author: evilsocket

CLI.md

@@ -50,7 +50,7 @@ $ dreadnode agent [OPTIONS] COMMAND [ARGS]...
 * `show`: Show the status of the active agent
 * `strikes`: List available strikes
 * `switch`: Switch to a different agent link
-* `templates`: List available agent templates with their...
+* `templates`: Interact with Strike templates
 * `versions`: List historical versions of the active agent
 
 ### `dreadnode agent clone`
@@ -108,7 +108,7 @@ $ dreadnode agent init [OPTIONS] STRIKE
 
 * `-d, --dir DIRECTORY`: The directory to initialize  [default: .]
 * `-n, --name TEXT`: The project name (used for container naming)
-* `-t, --template [rigging_basic|rigging_loop|nerve_basic]`: The template to use for the agent  [default: rigging_basic]
+* `-t, --template TEXT`: The template to use for the agent
 * `-s, --source TEXT`: Initialize the agent using a custom template from a github repository, ZIP archive URL or local folder
 * `-p, --path TEXT`: If --source has been provided, use --path to specify a subfolder to initialize from
 * `--help`: Show this message and exit.
@@ -256,12 +256,49 @@ $ dreadnode agent switch [OPTIONS] AGENT_OR_PROFILE [DIRECTORY]
 
 ### `dreadnode agent templates`
 
+Interact with Strike templates
+
+**Usage**:
+
+```console
+$ dreadnode agent templates [OPTIONS] COMMAND [ARGS]...
+```
+
+**Options**:
+
+* `--help`: Show this message and exit.
+
+**Commands**:
+
+* `install`: Install a template pack
+* `show`: List available agent templates with their...
+
+#### `dreadnode agent templates install`
+
+Install a template pack
+
+**Usage**:
+
+```console
+$ dreadnode agent templates install [OPTIONS] [SOURCE]
+```
+
+**Arguments**:
+
+* `[SOURCE]`: The source of the template pack  [default: dreadnode/basic-agents]
+
+**Options**:
+
+* `--help`: Show this message and exit.
+
+#### `dreadnode agent templates show`
+
 List available agent templates with their descriptions
 
 **Usage**:
 
 ```console
-$ dreadnode agent templates [OPTIONS]
+$ dreadnode agent templates show [OPTIONS]
 ```
 
 **Options**:

README.md

@@ -139,7 +139,10 @@ Interact with Strike agents:
 dreadnode agent strikes
 
 # list all available templates with their descriptions
-dreadnode agent templates
+dreadnode agent templates show
+
+# install a template pack from a github repository
+dreadnode agent templates install dreadnode/basic-templates
 
 # initialize a new agent in the current directory
 dreadnode agent init -t <template_name> <strike_id> 

dreadnode_cli/agent/cli.py

@@ -20,16 +20,19 @@
     format_run,
     format_runs,
     format_strikes,
-    format_templates,
 )
-from dreadnode_cli.agent.templates import Template, install_template, install_template_from_dir
+from dreadnode_cli.agent.templates import cli as templates_cli
+from dreadnode_cli.agent.templates.format import format_templates
+from dreadnode_cli.agent.templates.manager import TemplateManager
 from dreadnode_cli.config import UserConfig
 from dreadnode_cli.profile.cli import switch as switch_profile
 from dreadnode_cli.types import GithubRepo
-from dreadnode_cli.utils import download_and_unzip_archive, pretty_cli, repo_exists
+from dreadnode_cli.utils import download_and_unzip_archive, get_repo_archive_source_path, pretty_cli
 
 cli = typer.Typer(no_args_is_help=True)
 
+cli.add_typer(templates_cli, name="templates", help="Interact with Strike templates")
+
 
 def ensure_profile(agent_config: AgentConfig, *, user_config: UserConfig | None = None) -> None:
     """Ensure the active agent link matches the current server profile."""
@@ -66,26 +69,6 @@ def ensure_profile(agent_config: AgentConfig, *, user_config: UserConfig | None
         switch_profile(agent_config.active_link.profile)
 
 
-def get_repo_archive_source_path(source_dir: pathlib.Path) -> pathlib.Path:
-    """Return the actual source directory from a git repositoryZIP archive."""
-
-    if not (source_dir / "Dockerfile").exists() and not (source_dir / "Dockerfile.j2").exists():
-        # if src has been downloaded from a ZIP archive, it may contain a single
-        # '<user>-<repo>-<commit hash>' folder, that is the actual source we want to use.
-        # Check if source_dir contains only one folder and update it if so.
-        children = list(source_dir.iterdir())
-        if len(children) == 1 and children[0].is_dir():
-            source_dir = children[0]
-
-    return source_dir
-
-
-@cli.command(help="List available agent templates with their descriptions")
-@pretty_cli
-def templates() -> None:
-    print(format_templates())
-
-
 @cli.command(help="Initialize a new agent project")
 @pretty_cli
 def init(
@@ -98,8 +81,8 @@ def init(
         str | None, typer.Option("--name", "-n", help="The project name (used for container naming)")
     ] = None,
     template: t.Annotated[
-        Template, typer.Option("--template", "-t", help="The template to use for the agent")
-    ] = Template.rigging_basic,
+        str | None, typer.Option("--template", "-t", help="The template to use for the agent")
+    ] = None,
     source: t.Annotated[
         str | None,
         typer.Option(
@@ -137,18 +120,45 @@ def init(
     print(f":crossed_swords: Linking to strike '{strike_response.name}' ({strike_response.type})")
     print()
 
-    project_name = Prompt.ask("Project name?", default=name or directory.name)
+    project_name = Prompt.ask(":toolbox: Project name?", default=name or directory.name)
     print()
 
     directory.mkdir(exist_ok=True)
 
+    template_manager = TemplateManager()
     context = {"project_name": project_name, "strike": strike_response}
 
     if source is None:
-        # initialize from builtin template
-        template = Template(Prompt.ask("Template?", choices=[t.value for t in Template], default=template.value))
+        # get the templates that match the strike
+        available_templates = template_manager.get_templates_for_strike(strike_response.name, strike_response.type)
+        available: list[str] = list(available_templates.keys())
+
+        # none available
+        if not available:
+            if not template_manager.templates:
+                raise Exception(
+                    "No templates installed, use [bold]dreadnode agent templates install[/] to install some."
+                )
+            else:
+                raise Exception("No templates found for the given strike.")
+
+        # ask the user if the template has not been passed via command line
+        if template is None:
+            print(":notebook: Compatible templates:\n")
+            print(format_templates(available_templates, with_index=True))
+            print()
+
+            choice = Prompt.ask("Choice ", choices=[str(i + 1) for i in range(len(available))])
+            template = available[int(choice) - 1]
+
+        # validate the template
+        if template not in available:
+            raise Exception(
+                f"Template '{template}' not found, use [bold]dreadnode agent templates show[/] to see available templates."
+            )
 
-        install_template(template, directory, context)
+        # install the template
+        template_manager.install(template, directory, context)
     else:
         source_dir = pathlib.Path(source)
         cleanup = False
@@ -162,7 +172,7 @@ def init(
                 github_repo = GithubRepo(source)
 
                 # Check if the repo is accessible
-                if repo_exists(github_repo):
+                if github_repo.exists:
                     source_dir = download_and_unzip_archive(github_repo.zip_url)
 
                 # This could be a private repo that the user can access
@@ -193,7 +203,8 @@ def init(
             if path is not None:
                 source_dir = source_dir / path
 
-            install_template_from_dir(source_dir, directory, context)
+            # install the template
+            template_manager.install_from_dir(source_dir, directory, context)
         except Exception:
             if cleanup and source_dir.exists():
                 shutil.rmtree(source_dir)
@@ -521,7 +532,7 @@ def clone(
         shutil.rmtree(target)
 
     # Check if the repo is accessible
-    if repo_exists(github_repo):
+    if github_repo.exists:
         temp_dir = download_and_unzip_archive(github_repo.zip_url)
 
     # This could be a private repo that the user can access

dreadnode_cli/agent/format.py

@@ -9,7 +9,6 @@
 from rich.text import Text
 
 from dreadnode_cli import api
-from dreadnode_cli.agent.templates import Template, template_description
 
 P = t.ParamSpec("P")
 
@@ -312,14 +311,3 @@ def format_runs(runs: list[api.Client.StrikeRunSummaryResponse]) -> RenderableTy
         )
 
     return table
-
-
-def format_templates() -> RenderableType:
-    table = Table(box=box.ROUNDED)
-    table.add_column("template")
-    table.add_column("description")
-
-    for template in Template:
-        table.add_row(f"[bold magenta]{template.value}[/]", template_description(template))
-
-    return table

dreadnode_cli/agent/templates/__init__.py

@@ -1,76 +1,3 @@
-import enum
-import pathlib
-import typing as t
+from dreadnode_cli.agent.templates.cli import cli
 
-from jinja2 import Environment, FileSystemLoader
-from rich.prompt import Prompt
-
-TEMPLATES_DIR = pathlib.Path(__file__).parent.parent / "templates"
-
-
-class Template(str, enum.Enum):
-    rigging_basic = "rigging_basic"
-    rigging_loop = "rigging_loop"
-    nerve_basic = "nerve_basic"
-
-
-def template_description(template: Template) -> str:
-    """Return the description of a template."""
-
-    readme = TEMPLATES_DIR / template.value / "README.md"
-    if readme.exists():
-        return readme.read_text()
-
-    return ""
-
-
-def install_template(template: Template, dest: pathlib.Path, context: dict[str, t.Any]) -> None:
-    """Install a template into a directory."""
-    install_template_from_dir(TEMPLATES_DIR / template.value, dest, context)
-
-
-def install_template_from_dir(src: pathlib.Path, dest: pathlib.Path, context: dict[str, t.Any]) -> None:
-    """Install a template from a source directory into a destination directory."""
-
-    if not src.exists():
-        raise Exception(f"Source directory '{src}' does not exist")
-
-    elif not src.is_dir():
-        raise Exception(f"Source '{src}' is not a directory")
-
-    # check for Dockerfile in the directory
-    elif not (src / "Dockerfile").exists() and not (src / "Dockerfile.j2").exists():
-        raise Exception(f"Source directory {src} does not contain a Dockerfile")
-
-    env = Environment(loader=FileSystemLoader(src))
-
-    # iterate over all items in the source directory
-    for src_item in src.glob("**/*"):
-        # get the relative path of the item
-        src_item_path = str(src_item.relative_to(src))
-        # get the destination path
-        dest_item = dest / src_item_path
-
-        # if the destination item is not the root directory and it exists,
-        # ask the user if they want to overwrite it
-        if dest_item != dest and dest_item.exists():
-            if Prompt.ask(f":axe: Overwrite {dest_item}?", choices=["y", "n"], default="n") == "n":
-                continue
-
-        # if the source item is a file
-        if src_item.is_file():
-            # if the file has a .j2 extension, render it using Jinja2
-            if src_item.name.endswith(".j2"):
-                # we can read as text
-                content = src_item.read_text()
-                j2_template = env.get_template(src_item_path)
-                content = j2_template.render(context)
-                dest_item = dest / src_item_path.removesuffix(".j2")
-                dest_item.write_text(content)
-            else:
-                # otherwise, copy the file as is
-                dest_item.write_bytes(src_item.read_bytes())
-
-        # if the source item is a directory, create it in the destination
-        elif src_item.is_dir():
-            dest_item.mkdir(exist_ok=True)
+__all__ = ["cli"]