new: added --source argument to agent init (ENG-99)

.vscode/settings.json

@@ -6,5 +6,10 @@
             "source.organizeImports": "explicit"
         },
         "editor.defaultFormatter": "charliermarsh.ruff"
-    }
+    },
+    "python.testing.pytestArgs": [
+        "dreadnode_cli"
+    ],
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true
 }

CLI.md

@@ -88,7 +88,8 @@ $ 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]`: The template to use for the agent  [default: rigging_basic]
+* `-t, --template [rigging_basic|rigging_loop|nerve_basic]`: The template to use for the agent  [default: rigging_basic]
+* `-s, --source TEXT`: Initialize the agent using a custom template from a github repository, ZIP archive URL or local folder
 * `--help`: Show this message and exit.
 
 ### `dreadnode agent latest`

README.md

@@ -112,6 +112,15 @@ dreadnode agent init -t <template_name> <strike_id>
 # initialize a new agent in the specified directory
 dreadnode agent init -t <template_name> <strike_id> --dir <directory>
 
+# initialize a new agent using a custom template from a github repository
+dreadnode agent init -s username/repository <strike_id>
+
+# initialize a new agent using a custom template from a github branch/tag
+dreadnode agent init -s username/repository@custom-feature <strike_id>
+
+# initialize a new agent using a custom template from a ZIP archive URL
+dreadnode agent init -s https://example.com/template-archive.zip <strike_id>
+
 # push a new version of the agent
 dreadnode agent push
 

dreadnode_cli/agent/cli.py

@@ -1,4 +1,5 @@
 import pathlib
+import shutil
 import time
 import typing as t
 
@@ -21,9 +22,10 @@
     format_strikes,
     format_templates,
 )
-from dreadnode_cli.agent.templates import Template, install_template
+from dreadnode_cli.agent.templates import Template, install_template, install_template_from_dir
 from dreadnode_cli.config import UserConfig
-from dreadnode_cli.utils import pretty_cli
+from dreadnode_cli.types import GithubRepo
+from dreadnode_cli.utils import download_and_unzip_archive, pretty_cli, repo_exists
 
 cli = typer.Typer(no_args_is_help=True)
 
@@ -48,33 +50,92 @@ def init(
     template: t.Annotated[
         Template, typer.Option("--template", "-t", help="The template to use for the agent")
     ] = Template.rigging_basic,
+    source: t.Annotated[
+        str | None,
+        typer.Option(
+            "--source",
+            "-s",
+            help="Initialize the agent using a custom template from a github repository, ZIP archive URL or local folder",
+        ),
+    ] = None,
 ) -> None:
+    print(f":coffee: Fetching strike '{strike}' ...")
+
     client = api.create_client()
 
     try:
         strike_response = client.get_strike(strike)
     except Exception as e:
         raise Exception(f"Failed to find strike '{strike}': {e}") from e
 
-    print()
     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)
-    template = Template(Prompt.ask("Template?", choices=[t.value for t in Template], default=template.value))
+    print()
 
     directory.mkdir(exist_ok=True)
 
     try:
         AgentConfig.read(directory)
         if Prompt.ask(":axe: Agent config exists, overwrite?", choices=["y", "n"], default="n") == "n":
             return
+        print()
     except Exception:
         pass
 
-    AgentConfig(project_name=project_name, strike=strike).write(directory=directory)
+    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))
+
+        install_template(template, directory, context)
+    else:
+        source_dir = pathlib.Path(source)
+        cleanup = False
+
+        if not source_dir.exists():
+            # source is not a local folder, so it can be:
+            # - full ZIP archive URL
+            # - github compatible reference
+
+            try:
+                github_repo = GithubRepo(source)
+
+                # Check if the repo is accessible
+                if repo_exists(github_repo):
+                    source_dir = download_and_unzip_archive(github_repo.zip_url)
+
+                # This could be a private repo that the user can access
+                # by getting an access token from our API
+                elif github_repo.namespace == "dreadnode" and (
+                    github_access_token := client.get_github_access_token([github_repo.repo])
+                ):
+                    print(":key: Accessed private repository")
+                    source_dir = download_and_unzip_archive(
+                        github_repo.api_zip_url, headers={"Authorization": f"Bearer {github_access_token.token}"}
+                    )
 
-    install_template(template, directory, {"project_name": project_name, "strike": strike_response})
+                else:
+                    raise Exception(f"Repository '{github_repo}' not found or inaccessible")
+
+            except ValueError:
+                source_dir = download_and_unzip_archive(source)
+
+            # make sure the temporary directory is cleaned up
+            cleanup = True
+
+        try:
+            # initialize from local folder, validation performed inside install_template_from_dir
+            install_template_from_dir(source_dir, directory, context)
+        except Exception:
+            if cleanup and source_dir.exists():
+                shutil.rmtree(source_dir)
+            raise
+
+    # Wait to write this until after the template is installed
+    AgentConfig(project_name=project_name, strike=strike).write(directory=directory)
 
     print()
     print(f"Initialized [b]{directory}[/]")

dreadnode_cli/agent/templates/__init__.py

@@ -26,20 +26,59 @@ def template_description(template: Template) -> str:
 
 def install_template(template: Template, dest: pathlib.Path, context: dict[str, t.Any]) -> None:
     """Install a template into a directory."""
-    src = TEMPLATES_DIR / template.value
-    env = Environment(loader=FileSystemLoader(src))
+    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")
 
-    for src_item in src.iterdir():
-        dest_item = dest / src_item.name
-        content = src_item.read_text()
+    elif not (src / "Dockerfile").exists() and not (src / "Dockerfile.j2").exists():
+        # if src has been downloaded from a ZIP archive, it may contain a single
+        # 'project-main' folder, that is the actual source we want to use.
+        # Check if src contains only one folder and update it if so.
+        subdirs = [d for d in src.iterdir() if d.is_dir()]
+        if len(subdirs) == 1:
+            src = subdirs[0]
 
-        if src_item.name.endswith(".j2"):
-            j2_template = env.get_template(src_item.name)
-            content = j2_template.render(context)
-            dest_item = dest / src_item.name.removesuffix(".j2")
+        # check again for Dockerfile in the subdirectory
+        if not (src / "Dockerfile").exists() and not (src / "Dockerfile.j2").exists():
+            raise Exception("Source directory does not contain a Dockerfile")
 
-        if dest_item.exists():
-            if Prompt.ask(f":axe: Overwrite {dest_item.name}?", choices=["y", "n"], default="n") == "n":
+    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
 
-        dest_item.write_text(content)
+        # 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)

dreadnode_cli/agent/tests/test_templates.py

@@ -1,6 +1,8 @@
 import pathlib
 from unittest.mock import patch
 
+import pytest
+
 from dreadnode_cli.agent import templates
 
 
@@ -11,3 +13,139 @@ def test_templates_install(tmp_path: pathlib.Path) -> None:
     assert (tmp_path / "requirements.txt").exists()
     assert (tmp_path / "Dockerfile").exists()
     assert (tmp_path / "agent.py").exists()
+
+
+def test_templates_install_from_dir(tmp_path: pathlib.Path) -> None:
+    templates.install_template_from_dir(templates.TEMPLATES_DIR / "rigging_basic", tmp_path, {"name": "World"})
+
+    assert (tmp_path / "requirements.txt").exists()
+    assert (tmp_path / "Dockerfile").exists()
+    assert (tmp_path / "agent.py").exists()
+
+
+def test_templates_install_from_dir_with_dockerfile_template(tmp_path: pathlib.Path) -> None:
+    # create source directory
+    source_dir = tmp_path / "source"
+    source_dir.mkdir()
+
+    # create a Dockerfile.j2 template
+    dockerfile_content = """
+FROM python:3.9
+WORKDIR /app
+ENV APP_NAME={{name}}
+COPY . .
+CMD ["python", "app.py"]
+"""
+    (source_dir / "Dockerfile.j2").write_text(dockerfile_content)
+
+    # create destination directory
+    dest_dir = tmp_path / "dest"
+    dest_dir.mkdir()
+
+    # install template
+    templates.install_template_from_dir(source_dir, dest_dir, {"name": "TestContainer"})
+
+    # verify Dockerfile was rendered correctly
+    expected_dockerfile = """
+FROM python:3.9
+WORKDIR /app
+ENV APP_NAME=TestContainer
+COPY . .
+CMD ["python", "app.py"]
+"""
+    assert (dest_dir / "Dockerfile").exists()
+    assert (dest_dir / "Dockerfile").read_text().strip() == expected_dockerfile.strip()
+
+
+def test_templates_install_from_dir_nested_structure(tmp_path: pathlib.Path) -> None:
+    # create source directory with nested structure
+    source_dir = tmp_path / "source"
+    source_dir.mkdir()
+
+    # create some regular files
+    (source_dir / "Dockerfile").touch()
+    (source_dir / "README.md").write_text("# Test Project")
+
+    # create nested folders with files
+    config_dir = source_dir / "config"
+    config_dir.mkdir()
+    (config_dir / "settings.json").write_text('{"debug": true}')
+
+    templates_dir = source_dir / "templates"
+    templates_dir.mkdir()
+    (templates_dir / "base.html.j2").write_text("<html><body>Hello {{name}}!</body></html>")
+
+    src_dir = source_dir / "src"
+    src_dir.mkdir()
+    (src_dir / "main.py").touch()
+
+    # deeper nested folder
+    utils_dir = src_dir / "utils"
+    utils_dir.mkdir()
+    (utils_dir / "helpers.py").touch()
+    (utils_dir / "config.py.j2").write_text("APP_NAME = '{{name}}'")
+
+    # create destination directory
+    dest_dir = tmp_path / "dest"
+    dest_dir.mkdir()
+
+    # install template
+    templates.install_template_from_dir(source_dir, dest_dir, {"name": "TestApp"})
+
+    # verify regular files were copied
+    assert (dest_dir / "Dockerfile").exists()
+    assert (dest_dir / "README.md").read_text() == "# Test Project"
+
+    # verify nested structure and files
+    assert (dest_dir / "config" / "settings.json").read_text() == '{"debug": true}'
+    assert (dest_dir / "src" / "main.py").exists()
+    assert (dest_dir / "src" / "utils" / "helpers.py").exists()
+
+    # verify j2 templates were rendered correctly
+    assert (dest_dir / "templates" / "base.html").read_text() == "<html><body>Hello TestApp!</body></html>"
+    assert (dest_dir / "src" / "utils" / "config.py").read_text() == "APP_NAME = 'TestApp'"
+
+
+def test_templates_install_from_dir_missing_source(tmp_path: pathlib.Path) -> None:
+    source_dir = tmp_path / "nonexistent"
+    with pytest.raises(Exception, match="Source directory '.*' does not exist"):
+        templates.install_template_from_dir(source_dir, tmp_path, {"name": "World"})
+
+
+def test_templates_install_from_dir_source_is_file(tmp_path: pathlib.Path) -> None:
+    source_file = tmp_path / "source.txt"
+    source_file.touch()
+
+    with pytest.raises(Exception, match="Source '.*' is not a directory"):
+        templates.install_template_from_dir(source_file, tmp_path, {"name": "World"})
+
+
+def test_templates_install_from_dir_missing_dockerfile(tmp_path: pathlib.Path) -> None:
+    source_dir = tmp_path / "source"
+    source_dir.mkdir()
+    (source_dir / "agent.py").touch()
+
+    with pytest.raises(Exception, match="Source directory does not contain a Dockerfile"):
+        templates.install_template_from_dir(source_dir, tmp_path, {"name": "World"})
+
+
+def test_templates_install_from_dir_single_inner_folder(tmp_path: pathlib.Path) -> None:
+    # create a source directory with a single inner folder to simulate a github zip archive
+    source_dir = tmp_path / "source"
+    source_dir.mkdir()
+    inner_dir = source_dir / "project-main"
+    inner_dir.mkdir()
+
+    # create a Dockerfile in the inner directory
+    (inner_dir / "Dockerfile").touch()
+    (inner_dir / "agent.py").touch()
+
+    dest_dir = tmp_path / "dest"
+    dest_dir.mkdir()
+
+    # install from the outer directory - should detect and use inner directory
+    templates.install_template_from_dir(source_dir, dest_dir, {"name": "World"})
+
+    # assert files were copied from inner directory
+    assert (dest_dir / "Dockerfile").exists()
+    assert (dest_dir / "agent.py").exists()

dreadnode_cli/api.py

@@ -229,6 +229,19 @@ def submit_challenge_flag(self, challenge: str, flag: str) -> bool:
         response = self.request("POST", f"/api/challenges/{challenge}/submit-flag", json_data={"flag": flag})
         return bool(response.json().get("correct", False))
 
+    # Github
+
+    class GithubTokenResponse(BaseModel):
+        token: str
+        expires_at: datetime
+        repos: list[str]
+
+    def get_github_access_token(self, repos: list[str]) -> GithubTokenResponse | None:
+        """Try to get a GitHub access token for the given repositories."""
+
+        response = self.request("POST", "/api/github/token", json_data={"repos": repos})
+        return self.GithubTokenResponse(**response.json()) if response.status_code == 200 else None
+
     # Strikes
 
     StrikeRunStatus = t.Literal[