Add support for groups of secrets

packages/sshecret-admin/src/sshecret_admin/services/keepass.py

@@ -7,15 +7,16 @@ from pathlib import Path
 from typing import cast
 
 import pykeepass
-from .master_password import decrypt_master_password
 from sshecret_admin.core.settings import AdminServerSettings
 
+from .models import SecretGroup
+from .master_password import decrypt_master_password
+
 
 LOG = logging.getLogger(__name__)
 
 NO_USERNAME = "NO_USERNAME"
 
-
 DEFAULT_LOCATION = "keepass.kdbx"
 
 
@@ -25,6 +26,20 @@ def create_password_db(location: Path, password: str) -> None:
     pykeepass.create_database(str(location.absolute()), password=password)
 
 
+def _kp_group_to_secret_group(
+    kp_group: pykeepass.group.Group, parent: SecretGroup | None = None
+) -> SecretGroup:
+    """Convert keepass group to secret group dataclass."""
+    group_name = cast(str, kp_group.name)
+    group = SecretGroup(name=group_name, description=kp_group.notes)
+    if parent:
+        group.parent_group = parent
+    for subgroup in kp_group.subgroups:
+        group.children.append(_kp_group_to_secret_group(subgroup, group))
+
+    return group
+
+
 class PasswordContext:
     """Password Context class."""
 
@@ -32,15 +47,41 @@ class PasswordContext:
         """Initialize password context."""
         self.keepass: pykeepass.PyKeePass = keepass
 
-    def add_entry(self, entry_name: str, secret: str, overwrite: bool = False) -> None:
+    @property
+    def _root_group(self) -> pykeepass.group.Group:
+        """Return the root group."""
+        return cast(pykeepass.group.Group, self.keepass.root_group)
+
+    def _get_entry(self, name: str) -> pykeepass.entry.Entry | None:
+        """Get entry."""
+        entry = cast(
+            "pykeepass.entry.Entry | None",
+            self.keepass.find_entries(title=name, first=True),
+        )
+        return entry
+
+    def _get_group(self, name: str) -> pykeepass.group.Group | None:
+        """Find a group."""
+        group = cast(
+            pykeepass.group.Group | None,
+            self.keepass.find_groups(name=name, first=True),
+        )
+        return group
+
+    def add_entry(
+        self,
+        entry_name: str,
+        secret: str,
+        overwrite: bool = False,
+        group_name: str | None = None,
+    ) -> None:
         """Add an entry.
 
         Specify overwrite=True to overwrite the existing secret value, if it exists.
+        This will not move the entry, if the group_name is different from the original group.
+
         """
-        entry = cast(
-            "pykeepass.entry.Entry | None",
-            self.keepass.find_entries(title=entry_name, first=True),
-        )
+        entry = self._get_entry(entry_name)
         if entry and overwrite:
             entry.password = secret
             self.keepass.save()
@@ -48,9 +89,14 @@ class PasswordContext:
 
         if entry:
             raise ValueError("Error: A secret with this name already exists.")
-        LOG.debug("Add secret entry to keepass: %s", entry_name)
+        LOG.debug("Add secret entry to keepass: %s, group: %r", entry_name, group_name)
+        if group_name:
+            destination_group = self._get_group(group_name)
+        else:
+            destination_group = self._root_group
+
         entry = self.keepass.add_entry(
-            destination_group=self.keepass.root_group,
+            destination_group=destination_group,
             title=entry_name,
             username=NO_USERNAME,
             password=secret,
@@ -59,10 +105,7 @@ class PasswordContext:
 
     def get_secret(self, entry_name: str) -> str | None:
         """Get the secret value."""
-        entry = cast(
-            "pykeepass.entry.Entry | None",
-            self.keepass.find_entries(title=entry_name, first=True),
-        )
+        entry = self._get_entry(entry_name)
         if not entry:
             return None
 
@@ -72,9 +115,98 @@ class PasswordContext:
 
         raise RuntimeError(f"Cannot get password for entry {entry_name}")
 
-    def get_available_secrets(self) -> list[str]:
+    def get_secret_groups(self, pattern: str | None = None) -> list[SecretGroup]:
+        """Get secret groups.
+
+        A regex pattern may be provided to filter groups.
+        """
+        if pattern:
+            groups = cast(
+                list[pykeepass.group.Group],
+                self.keepass.find_groups(name=pattern, regex=True),
+            )
+        else:
+            all_groups = cast(list[pykeepass.group.Group], self.keepass.groups)
+            # We skip the root group
+            groups = [group for group in all_groups if not group.is_root_group]
+
+        secret_groups = [_kp_group_to_secret_group(group) for group in groups]
+        return secret_groups
+
+    def add_group(
+        self, name: str, description: str | None = None, parent_group: str | None = None
+    ) -> None:
+        """Add a group."""
+        kp_parent_group = self._root_group
+        if parent_group:
+            query = cast(
+                pykeepass.group.Group | None,
+                self.keepass.find_groups(name=parent_group, first=True),
+            )
+            if not query:
+                raise ValueError(
+                    f"Error: Cannot find a parent group named {parent_group}"
+                )
+            kp_parent_group = query
+        self.keepass.add_group(kp_parent_group, name, notes=description)
+        self.keepass.save()
+
+    def set_group_description(self, name: str, description: str) -> None:
+        """Set the description of a group."""
+        group = self._get_group(name)
+        if not group:
+            raise ValueError(f"Error: No such group {name}")
+
+        group.notes = description
+        self.keepass.save()
+
+    def set_secret_group(self, entry_name: str, group_name: str | None) -> None:
+        """Move a secret to a group.
+
+        If group is None, the secret will be placed in the root group.
+        """
+        entry = self._get_entry(entry_name)
+        if not entry:
+            raise ValueError(
+                f"Cannot find secret entry named {entry_name} in secrets database"
+            )
+        if group_name:
+            group = self._get_group(group_name)
+            if not group:
+                raise ValueError(f"Cannot find a group named {group_name}")
+        else:
+            group = self._root_group
+
+        self.keepass.move_entry(entry, group)
+        self.keepass.save()
+
+    def move_group(self, name: str, parent_group: str | None) -> None:
+        """Move a group.
+
+        If parent_group is None, it will be moved to the root.
+        """
+        group = self._get_group(name)
+        if not group:
+            raise ValueError(f"Error: No such group {name}")
+        if parent_group:
+            parent = self._get_group(parent_group)
+            if not parent:
+                raise ValueError(f"Error: No such group {parent_group}")
+        else:
+            parent = self._root_group
+
+        self.keepass.move_group(group, parent)
+        self.keepass.save()
+
+    def get_available_secrets(self, group_name: str | None = None) -> list[str]:
         """Get the names of all secrets in the database."""
-        entries = self.keepass.entries
+        if group_name:
+            group = self._get_group(group_name)
+            if not group:
+                raise ValueError(f"Error: No such group {group_name}")
+            entries = group.entries
+        else:
+            entries = cast(list[pykeepass.entry.Entry], self.keepass.entries)
         if not entries:
             return []
         return [str(entry.title) for entry in entries]
@@ -90,6 +222,28 @@ class PasswordContext:
         entry.delete()
         self.keepass.save()
 
+    def delete_group(self, name: str, keep_entries: bool = True) -> None:
+        """Delete a group.
+
+        If keep_entries is set to False, all entries in the group will be deleted.
+        """
+        group = self._get_group(name)
+        if not group:
+            return
+        if keep_entries:
+            for entry in cast(
+                list[pykeepass.entry.Entry],
+                self.keepass.find_entries(recursive=True, group=group),
+            ):
+                # Move the entry to the root group.
+                LOG.warning(
+                    "Moving orphaned secret entry %s to root group", entry.title
+                )
+                self.keepass.move_entry(entry, self._root_group)
+
+        self.keepass.delete_group(group)
+        self.keepass.save()
+
 
 @contextmanager
 def _password_context(location: Path, password: str) -> Iterator[PasswordContext]: