From 66e8222a09da0fbc27ada889acab28ba8374cf42 Mon Sep 17 00:00:00 2001
From: Maciej Borzecki <maciej.borzecki@canonical.com>
Date: Tue, 6 Feb 2024 14:55:04 +0100
Subject: [PATCH] subiquitycore: use a premade host key fingerprints info if
 present

It is possible that the platform integration glue may have already prepared a
summary of host key fingerprints at the state directory. If so, use it
otherwise, try to build the summary directly.

Signed-off-by: Maciej Borzecki <maciej.borzecki@canonical.com>
---
 subiquitycore/ssh.py            | 21 +++++++------
 subiquitycore/tests/test_ssh.py | 53 +++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 9 deletions(-)
 create mode 100644 subiquitycore/tests/test_ssh.py

diff --git a/subiquitycore/ssh.py b/subiquitycore/ssh.py
index 5fbbaf68..6a033b34 100644
--- a/subiquitycore/ssh.py
+++ b/subiquitycore/ssh.py
@@ -77,18 +77,21 @@ The {keytype} host key fingerprint is:
 )
 
 
-def host_key_info():
-    if os.getenv("SNAP_CONFINEMENT", "classic") == "strict":
-        # if we run in confinement, we have no direct accesss to host
-        # keys info use prepared finger prints if exist
-        snap_name = os.getenv("SNAP_NAME", "classic")
-        host_fingerprints = Path("/run/" + snap_name + "/host-fingerprints.txt")
+def host_key_info(runtime_state_dir=None):
+    if runtime_state_dir:
+        # host fingerprints information may have already been prepared by the
+        # platform glue
+        host_fingerprints = Path(runtime_state_dir) / "host-fingerprints.txt"
+        log.debug(
+            "pre-made host finterprints %s present: %s",
+            host_fingerprints,
+            host_fingerprints.is_file(),
+        )
         if host_fingerprints.is_file():
             fingerprints = open(host_fingerprints, "r")
             return fingerprints.read()
-        return ""
-    else:
-        return summarize_host_keys(host_key_fingerprints())
+
+    return summarize_host_keys(host_key_fingerprints())
 
 
 def summarize_host_keys(fingerprints):
diff --git a/subiquitycore/tests/test_ssh.py b/subiquitycore/tests/test_ssh.py
new file mode 100644
index 00000000..5cbf2f8c
--- /dev/null
+++ b/subiquitycore/tests/test_ssh.py
@@ -0,0 +1,53 @@
+# Copyright 2024 Canonical, Ltd.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+import os.path
+import tempfile
+import unittest
+from unittest.mock import patch
+
+from subiquitycore import ssh
+
+
+class TestSSH(unittest.TestCase):
+    @patch(
+        "subiquitycore.ssh.host_key_fingerprints",
+        return_value=[("key1-type", "key1-value"), ("key2-type", "key2-value")],
+    )
+    def test_host_key_info_premade(self, hkf):
+        # premade fingerprints are present
+        with tempfile.TemporaryDirectory(suffix="subiquity-ssh") as td:
+            fpfile = os.path.join(td, "host-fingerprints.txt")
+            with open(fpfile, "w") as outf:
+                outf.write("mock host fingerprints")
+
+            # fingerprints are pulled from the pre-made file
+            self.assertEqual(
+                ssh.host_key_info(runtime_state_dir=td), "mock host fingerprints"
+            )
+
+            # but are pulled from the system if the file is not there
+            os.remove(fpfile)
+            self.assertIn(
+                "key1-type key1-value", ssh.host_key_info(runtime_state_dir=td)
+            )
+
+    @patch(
+        "subiquitycore.ssh.host_key_fingerprints",
+        return_value=[("key1-type", "key1-value"), ("key2-type", "key2-value")],
+    )
+    def test_host_key_info_query(self, hkf):
+        self.assertIn("key1-type key1-value", ssh.host_key_info())
+        self.assertIn("key2-type key2-value", ssh.host_key_info())