diff --git a/stemcell_builder/stages/base_ssh/assets/10-ssh-firstboot-done.conf b/stemcell_builder/stages/base_ssh/assets/10-ssh-firstboot-done.conf
index 6101bf28ed..1d776524e8 100644
--- a/stemcell_builder/stages/base_ssh/assets/10-ssh-firstboot-done.conf
+++ b/stemcell_builder/stages/base_ssh/assets/10-ssh-firstboot-done.conf
@@ -1,2 +1,3 @@
 [Unit]
-ConditionPathExists=/root/firstboot_done
\ No newline at end of file
+Wants=firstboot.service
+After=firstboot.service
\ No newline at end of file
diff --git a/stemcell_builder/stages/base_ubuntu_firstboot/apply.sh b/stemcell_builder/stages/base_ubuntu_firstboot/apply.sh
index 37e0378211..63fff84e40 100755
--- a/stemcell_builder/stages/base_ubuntu_firstboot/apply.sh
+++ b/stemcell_builder/stages/base_ubuntu_firstboot/apply.sh
@@ -5,7 +5,12 @@ set -e
 base_dir=$(readlink -nf $(dirname $0)/../..)
 source $base_dir/lib/prelude_apply.bash
 
-cp $assets_dir/etc/rc.local $chroot/etc/rc.local
-cp $assets_dir/root/firstboot.sh $chroot/root/firstboot.sh
-chmod u+x "${chroot}/etc/rc.local"
-chmod 0755 $chroot/root/firstboot.sh
+install -D -m 0755 \
+  $assets_dir/root/firstboot.sh \
+  $chroot/root/firstboot.sh
+
+install -D -m 0644 \
+  $assets_dir/etc/systemd/system/firstboot.service \
+  $chroot/etc/systemd/system/firstboot.service
+
+run_in_chroot $chroot "systemctl enable firstboot.service"
diff --git a/stemcell_builder/stages/base_ubuntu_firstboot/assets/etc/rc.local b/stemcell_builder/stages/base_ubuntu_firstboot/assets/etc/rc.local
deleted file mode 100644
index dd5376a0dc..0000000000
--- a/stemcell_builder/stages/base_ubuntu_firstboot/assets/etc/rc.local
+++ /dev/null
@@ -1,22 +0,0 @@
-#!/bin/sh -e
-#execute firstboot.sh only once
-if [ ! -e /root/firstboot_done ]; then
-    if [ -e /root/firstboot.sh ]; then
-        MAX_RETRIES=5
-        COUNT=0
-        while [ $COUNT -lt $MAX_RETRIES ]; do
-            /root/firstboot.sh
-            if [ $? -eq 0 ]; then
-                touch /root/firstboot_done
-                break
-            fi
-            COUNT=$((COUNT+1))
-            sleep 1
-        done
-        if [ $COUNT -eq $MAX_RETRIES ]; then
-            echo "Max retries reached. Exiting..."
-            exit 1
-        fi
-    fi
-fi
-exit 0
\ No newline at end of file
diff --git a/stemcell_builder/stages/base_ubuntu_firstboot/assets/etc/systemd/system/firstboot.service b/stemcell_builder/stages/base_ubuntu_firstboot/assets/etc/systemd/system/firstboot.service
new file mode 100644
index 0000000000..cb2584e817
--- /dev/null
+++ b/stemcell_builder/stages/base_ubuntu_firstboot/assets/etc/systemd/system/firstboot.service
@@ -0,0 +1,12 @@
+[Unit]
+Description=Run first boot tasks
+ConditionFirstBoot=yes
+Before=ssh.service
+
+[Service]
+Type=oneshot
+ExecStart=/root/firstboot.sh
+RemainAfterExit=yes
+
+[Install]
+WantedBy=multi-user.target
diff --git a/stemcell_builder/stages/base_ubuntu_firstboot/assets/root/firstboot.sh b/stemcell_builder/stages/base_ubuntu_firstboot/assets/root/firstboot.sh
index ef85602b9c..365789753b 100755
--- a/stemcell_builder/stages/base_ubuntu_firstboot/assets/root/firstboot.sh
+++ b/stemcell_builder/stages/base_ubuntu_firstboot/assets/root/firstboot.sh
@@ -1,6 +1,7 @@
 #!/bin/sh
+set -e
 
 rm -f /etc/ssh/ssh_host*key*
+ssh-keygen -A -v
 
-dpkg-reconfigure -fnoninteractive -pcritical openssh-server
 dpkg-reconfigure -fnoninteractive sysstat
diff --git a/stemcell_builder/stages/system_azure_init/apply.sh b/stemcell_builder/stages/system_azure_init/apply.sh
index ac0f6f9a09..764c7b9a2c 100755
--- a/stemcell_builder/stages/system_azure_init/apply.sh
+++ b/stemcell_builder/stages/system_azure_init/apply.sh
@@ -46,12 +46,16 @@ cat > $chroot/etc/logrotate.d/waagent <<EOS
 }
 EOS
 
-#setup cloud-init
+# Setup cloud-init
 rm $chroot/etc/cloud/*.cfg
 rm $chroot/etc/cloud/cloud.cfg.d/*.cfg
 cp -f $dir/assets/etc/cloud-init/cloud.cfg $chroot/etc/cloud/cloud.cfg
 cp -f $dir/assets/etc/cloud-init/*-*.cfg $chroot/etc/cloud/cloud.cfg.d/
 
+# Ensures that cloud-init waits until host keys have been regenerated.
+mkdir -p $chroot/etc/systemd/system/cloud-config.service.d
+cp -f $dir/assets/etc/systemd/system/cloud-config.service.d/firstboot-blocker.conf \
+      $chroot/etc/systemd/system/cloud-config.service.d/firstboot-blocker.conf
 
 # this will append the following two relevant lines (plus a few commented out lines)
 # to the default-conf:
diff --git a/stemcell_builder/stages/system_azure_init/assets/etc/systemd/system/cloud-config.service.d/firstboot-blocker.conf b/stemcell_builder/stages/system_azure_init/assets/etc/systemd/system/cloud-config.service.d/firstboot-blocker.conf
new file mode 100644
index 0000000000..9c1b09bd0e
--- /dev/null
+++ b/stemcell_builder/stages/system_azure_init/assets/etc/systemd/system/cloud-config.service.d/firstboot-blocker.conf
@@ -0,0 +1,3 @@
+[Unit]
+Wants=firstboot.service
+After=firstboot.service
diff --git a/stemcell_builder/stages/system_azure_init/assets/etc/waagent/walinuxagent.service b/stemcell_builder/stages/system_azure_init/assets/etc/waagent/walinuxagent.service
index 55db8fa3d9..e4408d633a 100755
--- a/stemcell_builder/stages/system_azure_init/assets/etc/waagent/walinuxagent.service
+++ b/stemcell_builder/stages/system_azure_init/assets/etc/waagent/walinuxagent.service
@@ -7,17 +7,15 @@
 [Unit]
 Description=Azure Linux Agent
 
-After=network-online.target
-Wants=network-online.target ssh.service sshd-keygen.service
+# NON-DEFAULT: Must run after the firstboot.service, which regenerates the ssh keys
+After=firstboot.service network-online.target cloud-init.service
+Wants=network-online.target sshd.service sshd-keygen.service
 
 ConditionFileIsExecutable=/usr/sbin/waagent
 ConditionPathExists=/etc/waagent.conf
 
 [Service]
 Type=simple
-# stemcells on Azure re-generate the SSH Hostkey upon first reboot
-# waagent has to wait until the file was recreated
-ExecStartPre=/bin/bash -c "while [ ! -f /root/firstboot_done ]; do sleep 1; done"
 ExecStart=/usr/bin/python3 -u /usr/sbin/waagent -daemon
 Restart=always
 Slice=azure.slice