new: added --source-branch argument

CLI.md

@@ -90,6 +90,7 @@ $ dreadnode agent init [OPTIONS] STRIKE
 * `-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]
 * `--source TEXT`: Initialize the agent using a custom template from a github repository, ZIP archive URL or local folder
+* `--source-branch TEXT`: If --source is a github repository, use this as the branch name  [default: main]
 * `--help`: Show this message and exit.
 
 ### `dreadnode agent latest`

dreadnode_cli/agent/cli.py

@@ -56,6 +56,13 @@ def init(
             help="Initialize the agent using a custom template from a github repository, ZIP archive URL or local folder",
         ),
     ] = None,
+    source_branch: t.Annotated[
+        str,
+        typer.Option(
+            "--source-branch",
+            help="If --source is a github repository, use this as the branch name",
+        ),
+    ] = "main",
 ) -> None:
     print(f":coffee: Fetching strike '{strike}' ...")
 
@@ -100,7 +107,7 @@ def init(
             # - full github repository URL
             # - full ZIP archive URL
             # - username/repo
-            source = normalize_template_source(source)
+            source = normalize_template_source(source, source_branch)
             # download and unzip to a temporary directory
             source_dir = download_and_unzip_archive(source)
             # make sure the temporary directory is cleaned up

dreadnode_cli/tests/test_utils.py

@@ -104,6 +104,14 @@ def test_normalize_template_source_custom_zip_url() -> None:
     assert normalize_template_source(url) == url
 
 
+def test_normalize_template_source_custom_branch() -> None:
+    # Test with custom branch name
+    url = "user/repo"
+    branch = "develop"
+    expected = "https://github.com/user/repo/archive/refs/heads/develop.zip"
+    assert normalize_template_source(url, branch) == expected
+
+
 def test_download_and_unzip_archive_success(tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch) -> None:
     # create a mock zip file with test content
     test_file_content = b"test content"

dreadnode_cli/utils.py

@@ -67,7 +67,7 @@ def parse_jwt_token_expiration(token: str) -> datetime:
     return datetime.fromtimestamp(json.loads(payload).get("exp"))
 
 
-def normalize_template_source(source: str) -> str:
+def normalize_template_source(source: str, source_branch: str = "main") -> str:
     """Normalize a template source to a ZIP archive URL."""
 
     # github repository / ZIP archive URL
@@ -77,7 +77,7 @@ def normalize_template_source(source: str) -> str:
 
     if not source.lower().endswith(".zip"):
         # normalize to ZIP archive URL
-        source = f"{source}/archive/refs/heads/main.zip"
+        source = f"{source}/archive/refs/heads/{source_branch}.zip"
 
     return source