From c91979e520347a2185de9fb0cc2dae683c9491f3 Mon Sep 17 00:00:00 2001
From: Tomas Barton <barton.tomas@gmail.com>
Date: Thu, 16 Feb 2023 14:55:48 +0100
Subject: [PATCH] Support assigning multiple tags to a hostkey

---
 README.md                     | 16 +++++++++
 manifests/hostkeys.pp         | 27 +++++++++-----
 spec/classes/hostkeys_spec.rb | 67 +++++++++++++++++++++++++++++++++++
 3 files changed, 102 insertions(+), 8 deletions(-)
 create mode 100644 spec/classes/hostkeys_spec.rb

diff --git a/README.md b/README.md
index 3c6770ca..ae354185 100644
--- a/README.md
+++ b/README.md
@@ -348,6 +348,22 @@ class YOURCUSTOMCLASS {
 }
 ```
 
+## Tag hostkey
+
+Assign tags to exported `sshkey` resources (when `ssh::storeconfigs_enabled` is set to `true`).
+
+```yaml
+ssh::hostkeys::tags:
+  - hostkey_group1
+  - hostkey_group2
+```
+
+Host keys then can be imported using:
+
+```puppet
+Sshkey <<| tag == "hostkey_group1" |>>
+```
+
 ## Excluding network interfaces or ipaddresses
 
 Use hiera to exclude interfaces or ipaddresses from hostkey inclusion
diff --git a/manifests/hostkeys.pp b/manifests/hostkeys.pp
index 0d66c85e..28463d44 100644
--- a/manifests/hostkeys.pp
+++ b/manifests/hostkeys.pp
@@ -19,13 +19,17 @@
 # @param use_trusted_facts
 #   Whether to use trusted or normal facts
 #
+# @param tags
+#   Array of custom tags
+#
 class ssh::hostkeys (
-  Boolean             $export_ipaddresses  = true,
-  Optional[String[1]] $storeconfigs_group  = undef,
-  Array               $extra_aliases       = [],
-  Array               $exclude_interfaces  = [],
-  Array               $exclude_ipaddresses = [],
-  Boolean             $use_trusted_facts   = false,
+  Boolean                    $export_ipaddresses  = true,
+  Optional[String[1]]        $storeconfigs_group  = undef,
+  Array                      $extra_aliases       = [],
+  Array                      $exclude_interfaces  = [],
+  Array                      $exclude_ipaddresses = [],
+  Boolean                    $use_trusted_facts   = false,
+  Optional[Array[String[1]]] $tags                = undef,
 ) {
   if $use_trusted_facts {
     $fqdn_real = $trusted['certname']
@@ -44,8 +48,14 @@
     $host_aliases = sort(unique(flatten([$fqdn_real, $hostname_real, $extra_aliases])))
   }
 
-  if $storeconfigs_group {
-    tag 'hostkey_all', "hostkey_${storeconfigs_group}"
+  $storeconfigs_groups = $storeconfigs_group ? {
+    undef   => [],
+    default => ['hostkey_all', "hostkey_${storeconfigs_group}"],
+  }
+
+  $_tags = $tags ? {
+    undef   => $storeconfigs_groups,
+    default => $storeconfigs_groups + $tags,
   }
 
   ['dsa', 'rsa', 'ecdsa', 'ed25519'].each |String $key_type| {
@@ -63,6 +73,7 @@
         host_aliases => $host_aliases,
         type         => $key_type_real,
         key          => $facts['ssh'][$key_type]['key'],
+        tag          => $_tags,
       }
     } else {
       @@sshkey { "${fqdn_real}_${key_type}":
diff --git a/spec/classes/hostkeys_spec.rb b/spec/classes/hostkeys_spec.rb
new file mode 100644
index 00000000..42a64859
--- /dev/null
+++ b/spec/classes/hostkeys_spec.rb
@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe 'ssh::hostkeys', type: 'class' do
+  _, os_facts = on_supported_os.first
+
+  let(:facts) { os_facts }
+
+  context 'with tags' do
+    let(:params) do
+      {
+        tags: ['group1', 'group2']
+      }
+    end
+
+    ['dsa', 'rsa', 'ecdsa', 'ed25519'].each do |key_type|
+      it {
+        expect(exported_resources).to contain_sshkey("foo.example.com_#{key_type}")
+          .with(
+            ensure: 'present',
+            type: %r{^#{key_type}},
+            tag: ['group1', 'group2'],
+          )
+      }
+    end
+  end
+
+  context 'with storeconfigs_group' do
+    let(:params) do
+      {
+        storeconfigs_group: 'server_group',
+      }
+    end
+
+    ['dsa', 'rsa', 'ecdsa', 'ed25519'].each do |key_type|
+      it {
+        expect(exported_resources).to contain_sshkey("foo.example.com_#{key_type}")
+          .with(
+            ensure: 'present',
+            type: %r{^#{key_type}},
+            tag: ['hostkey_all', 'hostkey_server_group'],
+          )
+      }
+    end
+  end
+
+  context 'with storeconfigs_group and tags' do
+    let(:params) do
+      {
+        storeconfigs_group: 'server_group',
+        tags: ['group1', 'group2'],
+      }
+    end
+
+    ['dsa', 'rsa', 'ecdsa', 'ed25519'].each do |key_type|
+      it {
+        expect(exported_resources).to contain_sshkey("foo.example.com_#{key_type}")
+          .with(
+            ensure: 'present',
+            type: %r{^#{key_type}},
+            tag: ['hostkey_all', 'hostkey_server_group', 'group1', 'group2'],
+          )
+      }
+    end
+  end
+end