From 13e8eed592ffb6b2ae9a6b68f27e31c6cfd5ade6 Mon Sep 17 00:00:00 2001
From: Gerd Naschenweng <magicdude4eva@users.noreply.github.com>
Date: Sun, 18 Dec 2022 11:49:06 +0100
Subject: [PATCH 1/4] Updated compatibility for new UDM firmwares

---
 install.sh | 26 ++++++++++++++++++++++----
 1 file changed, 22 insertions(+), 4 deletions(-)

diff --git a/install.sh b/install.sh
index 6407a58..9f621be 100644
--- a/install.sh
+++ b/install.sh
@@ -1,15 +1,33 @@
 #!/bin/sh
 
-user_authorized_keys_file="/mnt/data/ssh/authorized_keys"
+# Get DataDir location
+DATA_DIR="/mnt/data"
+case "$(ubnt-device-info firmware || true)" in
+    1*)
+      DATA_DIR="/mnt/data"
+      ;;
+    2*)
+      DATA_DIR="/data"
+      ;;
+    3*)
+      DATA_DIR="/data"
+      ;;
+    *)
+      echo "ERROR: No persistent storage found." 1>&2
+      exit 1
+      ;;
+  esac
+
+user_authorized_keys_file="$DATA_DIR/ssh/authorized_keys"
 
 curl -sO https://raw.githubusercontent.com/fire1ce/UDM-Persistent-SSH-Keys/main/99-ssh-keys.sh
-mv 99-ssh-keys.sh /mnt/data/on_boot.d/99-ssh-keys.sh
-chmod +x /mnt/data/on_boot.d/99-ssh-keys.sh
+mv 99-ssh-keys.sh $DATA_DIR/on_boot.d/99-ssh-keys.sh
+chmod +x $DATA_DIR/on_boot.d/99-ssh-keys.sh
 echo "99-ssh-keys.sh installed"
 
 if [ ! -f "$user_authorized_keys_file" ]; then
   echo "Creating $user_authorized_keys_file"
-  mkdir -p /mnt/data/ssh
+  mkdir -p $DATA_DIR/ssh
   touch "$user_authorized_keys_file"
   chmod 0644 "$user_authorized_keys_file"
 fi

From 412b4c1af343f46e32e8f8a5c14c87dac868e4e8 Mon Sep 17 00:00:00 2001
From: Gerd Naschenweng <magicdude4eva@users.noreply.github.com>
Date: Sun, 18 Dec 2022 11:52:09 +0100
Subject: [PATCH 2/4] Improved support for new UDM firmwares

---
 install.sh | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/install.sh b/install.sh
index 9f621be..842e70c 100644
--- a/install.sh
+++ b/install.sh
@@ -19,6 +19,7 @@ case "$(ubnt-device-info firmware || true)" in
   esac
 
 user_authorized_keys_file="$DATA_DIR/ssh/authorized_keys"
+udm_authorized_keys_file="/root/.ssh/authorized_keys"
 
 curl -sO https://raw.githubusercontent.com/fire1ce/UDM-Persistent-SSH-Keys/main/99-ssh-keys.sh
 mv 99-ssh-keys.sh $DATA_DIR/on_boot.d/99-ssh-keys.sh
@@ -28,8 +29,9 @@ echo "99-ssh-keys.sh installed"
 if [ ! -f "$user_authorized_keys_file" ]; then
   echo "Creating $user_authorized_keys_file"
   mkdir -p $DATA_DIR/ssh
-  touch "$user_authorized_keys_file"
-  chmod 0644 "$user_authorized_keys_file"
+  mkdir -p /root/.ssh/
+  touch "$user_authorized_keys_file $udm_authorized_keys_file"
+  chmod 0644 "$user_authorized_keys_file $udm_authorized_keys_file"
 fi
 
 echo "==> Add you public keys to $user_authorized_keys_file"

From 5d278448cc8261176bd711fcf9f3d02e3ed2075f Mon Sep 17 00:00:00 2001
From: Gerd Naschenweng <magicdude4eva@users.noreply.github.com>
Date: Sun, 18 Dec 2022 12:02:11 +0100
Subject: [PATCH 3/4] Adjusted data directories depending on firmware

---
 README.md | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 84adf52..fbffbe7 100644
--- a/README.md
+++ b/README.md
@@ -10,10 +10,13 @@ This can be accomplished with a boot script. Flow this guide: [UDM / UDMPro Boot
 - Tested on [UDM PRO][amz-udm-pro-url]
 - UDM Pro doesn't support **ed25519** SSH Keys
 
+
 ## Installation
 
 The script was tested on UDM PRO
 
+(!) Depending on firmware your $DATA_DIR will be "/mnt/data" (Firmware 1.x) or "/data" (Firmware 2.x and 3.x)
+
 ```shell
 curl -s https://raw.githubusercontent.com/fire1ce/UDM-Persistent-SSH-Keys/main/install.sh | sh
 ```
@@ -21,7 +24,7 @@ curl -s https://raw.githubusercontent.com/fire1ce/UDM-Persistent-SSH-Keys/main/i
 Add you public RSA keys to:
 
 ```shell
-/mnt/data/ssh/authorized_keys
+$DATA_DIR/ssh/authorized_keys
 ```
 
 ## Uninstall
@@ -29,23 +32,23 @@ Add you public RSA keys to:
 Delete the 99-ssh-keys.sh file
 
 ```shell
-rm -rf /mnt/data/on_boot.d/99-ssh-keys.sh
+rm -rf $DATA_DIR/on_boot.d/99-ssh-keys.sh
 ```
 
 Delete your authorized_keys file
 
 ```shell
-rm -rf /mnt/data/ssh/authorized_keys
+rm -rf $DATA_DIR/ssh/authorized_keys
 ```
 
 ## Usage
 
-At boot the script with read the **/mnt/data/ssh/authorized_keys** file and add the content to UDM's **/root/.ssh/authorized_keys**
+At boot the script with read the **$DATA_DIR/ssh/authorized_keys** file and add the content to UDM's **/root/.ssh/authorized_keys**
 
 Manual run:
 
 ```shell
-/mnt/data/on_boot.d/99-ssh-keys.sh
+$DATA_DIR/on_boot.d/99-ssh-keys.sh
 ```
 
 <!-- --- -->

From 30b0dabf9ed9226cbaaa3df077116a3516b938ac Mon Sep 17 00:00:00 2001
From: Gerd Naschenweng <magicdude4eva@users.noreply.github.com>
Date: Sun, 18 Dec 2022 12:02:47 +0100
Subject: [PATCH 4/4] Update README.md

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index fbffbe7..0b09953 100644
--- a/README.md
+++ b/README.md
@@ -15,7 +15,7 @@ This can be accomplished with a boot script. Flow this guide: [UDM / UDMPro Boot
 
 The script was tested on UDM PRO
 
-(!) Depending on firmware your $DATA_DIR will be "/mnt/data" (Firmware 1.x) or "/data" (Firmware 2.x and 3.x)
+(!) Depending on firmware your `$DATA_DIR` will be `/mnt/data` (Firmware 1.x) or `/data` (Firmware 2.x and 3.x)
 
 ```shell
 curl -s https://raw.githubusercontent.com/fire1ce/UDM-Persistent-SSH-Keys/main/install.sh | sh