[backport/1.4] core: move setup_ssh() over to startup_blueos_update

core/tools/blueos_startup_update/blueos_startup_update.py

@@ -5,6 +5,8 @@
 import os
 import re
 import time
+import subprocess
+from pathlib import Path
 from typing import List, Tuple
 import configparser
 
@@ -664,6 +666,47 @@ def update_swap_size() -> bool:
         return False
 
 
+def setup_ssh() -> bool:
+    logger.info("Setting up SSH...")
+    # store the key in the docker .config volume
+    key_path = Path("/root/.config/.ssh")
+    private_key = key_path / "id_rsa"
+    public_key = private_key.with_suffix(".pub")
+    user = os.environ.get("SSH_USER", "pi")
+    gid = int(os.environ.get("USER_GID", 1000))
+    uid = int(os.environ.get("USER_UID", 1000))
+    authorized_keys = Path(f"/home/{user}/.ssh/authorized_keys")
+
+    try:
+        key_path.mkdir(parents=True, exist_ok=True)
+        # check if id_rsa.pub exists, creates a new one if it doesnt
+        if not public_key.is_file():
+            subprocess.run(["ssh-keygen", "-t", "rsa", "-f", private_key, "-q", "-N", ""], check=True)
+            logger.info("Generated new SSH key pair")
+
+        public_key_text = public_key.read_text("utf-8")
+        # add id_rsa.pub to authorized_keys if not there already
+        try:
+            authorized_keys_text = authorized_keys.read_text("utf-8")
+        except FileNotFoundError:
+            logger.info(f"File does not exist: {authorized_keys}")
+            authorized_keys_text = ""
+
+        if public_key_text not in authorized_keys_text:
+            if not authorized_keys_text.endswith("\n"):
+                authorized_keys_text += "\n"
+            authorized_keys_text += public_key_text
+            authorized_keys.write_text(authorized_keys_text, "utf-8")
+            logger.info("Added public key to authorized_keys")
+
+        os.chown(authorized_keys, uid, gid)
+        authorized_keys.chmod(0o600)
+        return True
+    except Exception as sshError:
+        logger.error(f"Error setting up ssh: {sshError}")
+        return False
+
+
 def main() -> int:
     start = time.time()
     # check if boot_loop_detector exists
@@ -747,6 +790,10 @@ def main() -> int:
 
 
 if __name__ == "__main__":
+    try:
+        setup_ssh()
+    except Exception as error:
+        logger.error(f"An error occurred while setting up ssh: {error}")
     try:
         main()
         if os.path.exists(BOOT_LOOP_DETECTOR):