Additional SSH key pair Algorithms

HISTORY.md

@@ -1,5 +1,10 @@
 # Release History
 
+## 0.5.3 (2024-12-12)
+
+- Update to to support more SSH key pair algorithms. (See https://docs.gitlab.com/ee/user/ssh.html#supported-ssh-key-types)
+- Modernize to use `Pathlib` instead of `os.path`.
+
 ## 0.5.2 (2024-11-12)
 
 - Update to pygit2 usage to enable usage of newer versions of the package while also maintaining backwards compatibility.

README.md

@@ -109,7 +109,7 @@ pytest
 To run linting:
 
 ```bash
-flake8 --config flake8.cfg battenberg
+flake8 --config flake8.cfg .
 ```
 
 ## Releasing a new version to PyPI

battenberg/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '0.5.2'
+__version__ = '0.5.3'
 
 
 from battenberg.core import Battenberg

battenberg/cli.py

@@ -1,15 +1,15 @@
-import os
 import sys
 import logging
 import subprocess
 from typing import Optional
 import click
+from pathlib import Path
 
 from battenberg.core import Battenberg
 from battenberg.utils import open_repository, open_or_init_repository
 from battenberg.errors import MergeConflictException
 
-sys.path.append(os.path.join(os.path.dirname(__file__), '..'))  # noqa: E402
+sys.path.append(Path(__file__).parent / '..')  # noqa: E402
 
 
 logging.basicConfig(level=logging.INFO)
@@ -23,9 +23,9 @@
 @click.group()
 @click.option(
     '-O',
-    default='.',
+    default=Path(),
     help='Direct the output of battenberg to this path instead of the current working directory.',
-    type=click.Path()
+    type=click.Path(path_type=Path)
 )
 @click.option(
     '--verbose',
@@ -94,9 +94,9 @@ def install(ctx, template: str, initial_branch: Optional[str], **kwargs):
 )
 @click.option(
     '--context-file',
-    default='.cookiecutter.json',
+    default=Path('.cookiecutter.json'),
     help='Path where we can find the output of the cookiecutter template context',
-    type=click.Path()
+    type=click.Path(path_type=Path)
 )
 @click.option(
     '--no-input',

battenberg/core.py

@@ -1,10 +1,9 @@
 import sys
-import os
 import json
 import logging
-import shutil
 import tempfile
 from typing import Any, Dict, Optional
+from pathlib import Path
 
 from pygit2 import (
     RemoteCallbacks,
@@ -55,6 +54,7 @@ def _fetch_remote_template(self):
 
     def _cookiecut(self, cookiecutter_kwargs: dict, worktree: TemporaryWorktree):
         with tempfile.TemporaryDirectory() as tmpdir:
+            tmpdir = Path(tmpdir)
             logger.debug(f"Cookiecutting {cookiecutter_kwargs['template']} into {tmpdir}")
             try:
                 cookiecutter(
@@ -70,12 +70,12 @@ def _cookiecut(self, cookiecutter_kwargs: dict, worktree: TemporaryWorktree):
 
             # Cookiecutter guarantees a single top-level directory after templating.
             logger.debug('Shifting directories down a level')
-            top_level_dir = os.path.join(tmpdir, os.listdir(tmpdir)[0])
-            for f in os.listdir(top_level_dir):
-                shutil.move(os.path.join(top_level_dir, f), worktree.path)
+            top_level_dir = tmpdir.joinpath(next(tmpdir.iterdir()))
+            for f in top_level_dir.iterdir():
+                f.rename(worktree.path / f.name)
 
-    def _get_context(self, context_file: str, base_path: str = None) -> Dict[str, Any]:
-        with open(os.path.join(base_path or self.repo.workdir, context_file)) as f:
+    def _get_context(self, context_file: Path, base_path: Path = None) -> Dict[str, Any]:
+        with (base_path or Path(self.repo.workdir)).joinpath(context_file).open() as f:
             return json.load(f)
 
     def _merge_template_branch(self, message: str, merge_target: str = None):
@@ -206,7 +206,8 @@ def install(self, template: str, checkout: Optional[str] = None, no_input: bool
         self._merge_template_branch(f'Installed template \'{template}\'')
 
     def upgrade(self, checkout: Optional[str] = None, no_input: bool = True,
-                merge_target: Optional[str] = None, context_file: str = '.cookiecutter.json'):
+                merge_target: Optional[str] = None,
+                context_file: Path = Path('.cookiecutter.json')):
         """Updates a repo using the found template context.
 
         Generates and applies any updates from the current repo state to the template state defined
@@ -275,7 +276,8 @@ def upgrade(self, checkout: Optional[str] = None, no_input: bool = True,
             )
             commit = worktree.repo.get(oid)
 
-        # Make template branch ref to created commit, see https://github.com/libgit2/pygit2/blob/master/CHANGELOG.md#1150-2024-05-18
+        # Make template branch ref to created commit,
+        # see https://github.com/libgit2/pygit2/blob/master/CHANGELOG.md#1150-2024-05-18
         self.repo.lookup_branch(TEMPLATE_BRANCH).set_target(str(commit.id))
 
         # Let's merge our changes into HEAD

battenberg/errors.py

@@ -1,3 +1,6 @@
+from pathlib import Path
+
+
 class BattenbergException(Exception):
     """
     Abstract Battenberg generic exception.
@@ -19,7 +22,7 @@ class WorktreeException(BattenbergException):
     Error raised when worktree could not be initialized.
     """
 
-    def __init__(self, worktree_name: str, worktree_path: str):
+    def __init__(self, worktree_name: str, worktree_path: Path):
         super().__init__(
             f'Worktree \'{worktree_name}\' could not be initialized in path \'{worktree_path}\''
         )
@@ -64,5 +67,12 @@ class InvalidRepositoryException(BattenbergException):
     Error raised when Git repository is invalid.
     """
 
-    def __init__(self, path: str):
+    def __init__(self, path: Path):
         super().__init__(f'{path} is not a valid repository path.')
+
+
+class KeypairException(BattenbergException):
+    """
+    Error raised when keypair could not be created.
+    """
+    pass

battenberg/temporary_worktree.py

@@ -1,10 +1,9 @@
-import os
 import shutil
 import tempfile
 import logging
 from types import TracebackType
 from typing import Optional, Type
-
+from pathlib import Path
 from pygit2 import Repository, Worktree
 from battenberg.errors import (
     RepositoryEmptyException,
@@ -25,8 +24,8 @@ def __init__(self, upstream: Repository, name: str, empty: bool = True):
         self.upstream = upstream
         self.name = name
         # Create the worktree working directory in the /tmp directory so it's out of the way.
-        self.tmp = tempfile.mkdtemp()
-        self.path = os.path.join(self.tmp, name)
+        self.tmp = Path(tempfile.mkdtemp())
+        self.path = self.tmp.joinpath(name)
         self.worktree = None
         self.repo = None
         self.empty = empty
@@ -48,10 +47,10 @@ def __enter__(self) -> 'TemporaryWorktree':
 
         if self.empty:
             for entry in self.repo[self.repo.head.target].tree:
-                if os.path.isdir(os.path.join(self.path, entry.name)):
-                    shutil.rmtree(os.path.join(self.path, entry.name))
+                if self.path.joinpath(entry.name).is_dir():
+                    shutil.rmtree(self.path.joinpath(entry.name))
                 else:
-                    os.remove(os.path.join(self.path, entry.name))
+                    self.path.joinpath(entry.name).unlink()
 
         logger.debug(f'Successfully created temporary worktree at {self.path}.')
 

battenberg/utils.py

@@ -1,16 +1,15 @@
-import os
 import logging
 import re
 import subprocess
 from typing import Optional
 from pygit2 import discover_repository, init_repository, Keypair, Repository
-from battenberg.errors import InvalidRepositoryException
-
+from battenberg.errors import InvalidRepositoryException, KeypairException
+from pathlib import Path
 
 logger = logging.getLogger(__name__)
 
 
-def open_repository(path: str) -> Repository:
+def open_repository(path: Path) -> Repository:
     try:
         repo_path = discover_repository(path)
     except Exception as e:
@@ -26,7 +25,7 @@ def open_repository(path: str) -> Repository:
     return Repository(repo_path)
 
 
-def open_or_init_repository(path: str, template: str, initial_branch: Optional[str] = None):
+def open_or_init_repository(path: Path, template: str, initial_branch: Optional[str] = None):
     try:
         return open_repository(path)
     except InvalidRepositoryException:
@@ -64,11 +63,38 @@ def set_initial_branch(repo: Repository, template: str):
             repo.references['HEAD'].set_target(initial_branch)
 
 
-def construct_keypair(public_key_path: str = None, private_key_path: str = None,
-                      passphrase: str = '') -> Keypair:
-    ssh_path = os.path.join(os.path.expanduser('~'), '.ssh')
-    if not public_key_path:
-        public_key_path = os.path.join(ssh_path, 'id_rsa.pub')
-    if not private_key_path:
-        private_key_path = os.path.join(ssh_path, 'id_rsa')
+ALGORITHMS = {
+    "ED25519": {
+        "public": "id_ed25519.pub",
+        "private": "id_ed25519"
+    },
+    "ED25519_SK": {
+        "public": "id_ed25519_sk.pub",
+        "private": "id_ed25519_sk"
+    },
+    "ECDSA_SK": {
+        "public": "id_ecdsa_sk.pub",
+        "private": "id_ecdsa_sk"
+    },
+    "RSA": {
+        "public": "id_rsa.pub",
+        "private": "id_rsa"
+    }
+}
+
+
+def find_key_paths(ssh_path: Path):
+    for algorithm in ALGORITHMS.values():
+        public_key_path = ssh_path / algorithm['public']
+        private_key_path = ssh_path / algorithm['private']
+        if public_key_path.exists() and private_key_path.exists():
+            return public_key_path, private_key_path
+
+    raise KeypairException(f'Could not find keypair in {ssh_path}',
+                           f'Possible options include: {ALGORITHMS}')
+
+
+def construct_keypair(ssh_path: Path = None, passphrase: str = '') -> Keypair:
+    ssh_path = ssh_path or Path('~/.ssh').expanduser()
+    public_key_path, private_key_path = find_key_paths(ssh_path)
     return Keypair("git", public_key_path, private_key_path, passphrase)

setup.py

@@ -1,15 +1,14 @@
-import io
 import re
 from setuptools import setup
+from pathlib import Path
 
-
-with open('README.md') as readme_file:
+with Path('README.md').open() as readme_file:
     readme = readme_file.read()
 
-with open('HISTORY.md') as history_file:
+with Path('HISTORY.md').open() as history_file:
     history = history_file.read()
 
-with io.open('battenberg/__init__.py', 'rt', encoding='utf8') as f:
+with Path('battenberg/__init__.py').open('rt', encoding='utf8') as f:
     version = re.search(r'__version__ = \'(.*?)\'', f.read()).group(1)
 
 

tests/conftest.py

@@ -1,13 +1,12 @@
-import os
-import stat
-from distutils.dir_util import copy_tree, remove_tree
+import shutil
 import tarfile
 import tempfile
 from types import TracebackType
 from typing import List, Optional, Type
 import pytest
 from pygit2 import Repository, init_repository
 from battenberg.core import Battenberg
+from pathlib import Path
 
 
 class TemporaryRepository:
@@ -19,18 +18,18 @@ class TemporaryRepository:
 
     def __enter__(self) -> 'TemporaryRepository':
         name = 'testrepo'
-        repo_path = os.path.join(os.path.dirname(__file__), 'data', f'{name}.tar')
-        self.temp_dir = tempfile.mkdtemp()
-        temp_repo_path = os.path.join(self.temp_dir, name)
+        repo_path = Path(__file__).parent / 'data' / f'{name}.tar'
+        self.temp_dir = Path(tempfile.mkdtemp())
+        temp_repo_path = self.temp_dir / name
         tar = tarfile.open(repo_path)
         tar.extractall(self.temp_dir)
         tar.close()
         return temp_repo_path
 
     def __exit__(self, type: Optional[Type[BaseException]], value: Optional[BaseException],
                  traceback: TracebackType):
-        if os.path.exists(self.temp_dir):
-           remove_tree(self.temp_dir)
+        if self.temp_dir.exists():
+            shutil.rmtree(self.temp_dir)
 
 
 @pytest.fixture
@@ -39,11 +38,10 @@ def repo() -> Repository:
         yield Repository(repo_path)
 
 
-def copy_template(repo: Repository, name : str, commit_message: str, parents: List[str] = None) -> str:
-    template_path = os.path.join(os.path.dirname(__file__), 'data', name)
-    # Use distuils implementation instead of shutil to allow for copying into
-    # a destination with existing files. See: https://stackoverflow.com/a/31039095/724251
-    copy_tree(template_path, repo.workdir)
+def copy_template(repo: Repository, name: str, commit_message: str,
+                  parents: List[str] = None) -> str:
+    template_path = Path(__file__).parent / 'data' / name
+    shutil.copytree(template_path, repo.workdir, dirs_exist_ok=True)
 
     # Stage the template changes
     repo.index.add_all()

tests/test_cli.py

@@ -1,11 +1,11 @@
+from pathlib import Path
 from typing import Dict
 from unittest.mock import Mock, patch
 import pytest
 from click.testing import CliRunner
-from cookiecutter.exceptions import CookiecutterException
 from pygit2 import Repository
 from battenberg import cli
-from battenberg.errors import BattenbergException, MergeConflictException
+from battenberg.errors import MergeConflictException
 
 
 @pytest.fixture
@@ -56,7 +56,7 @@ def test_upgrade(Battenberg: Mock, obj: Dict):
     assert result.output == ''
     Battenberg.return_value.upgrade.assert_called_once_with(
         checkout=None,
-        context_file='.cookiecutter.json',
+        context_file=Path('.cookiecutter.json'),
         merge_target=None,
         no_input=False
     )

tests/test_core.py

@@ -43,7 +43,8 @@ def test_install_raises_template_conflict(repo: Repository, template_repo: Repos
 
 
 @patch('battenberg.core.cookiecutter')
-def test_install_raises_failed_hook(cookiecutter: Mock, repo: Repository, template_repo: Repository):
+def test_install_raises_failed_hook(cookiecutter: Mock, repo: Repository,
+                                    template_repo: Repository):
     cookiecutter.side_effect = FailedHookException
 
     battenberg = Battenberg(repo)
@@ -105,7 +106,7 @@ def test_update_merge_target(installed_repo: Repository, template_repo: Reposito
 
     template_upgrade_message = f'commit (merge): Upgraded template \'{template_repo.workdir}\''
     main_merge_ref = find_ref_from_message(installed_repo, template_upgrade_message,
-                                             ref_name=merge_target)
+                                           ref_name=merge_target)
     assert main_merge_ref
     # Ensure the merge commit on the merge target branch was derived from the template branch.
     assert template_upgrade_oid in set(installed_repo[main_merge_ref.oid_new].parent_ids)