diff --git a/images/linux/scripts/base/apt.sh b/images/linux/scripts/base/apt.sh
index 1efbb6721..cf221452c 100644
--- a/images/linux/scripts/base/apt.sh
+++ b/images/linux/scripts/base/apt.sh
@@ -1,9 +1,5 @@
 #!/bin/bash -e
 
-export DEBIAN_FRONTEND=noninteractive
-apt-get -yq update
-apt-get -yq dist-upgrade
-
 # Stop and disable apt-daily upgrade services;
 systemctl stop apt-daily.timer
 systemctl disable apt-daily.timer
@@ -16,7 +12,7 @@ systemctl disable apt-daily-upgrade.service
 sudo sed -i 's/APT::Periodic::Update-Package-Lists "1"/APT::Periodic::Update-Package-Lists "0"/' /etc/apt/apt.conf.d/20auto-upgrades
 
 # Enable retry logic for apt up to 10 times
-echo "APT::Acquire::Retries \"10\";" > /etc/apt/apt.conf.d/80-retries 
+echo "APT::Acquire::Retries \"10\";" > /etc/apt/apt.conf.d/80-retries
 
 # Configure apt to always assume Y
 echo "APT::Get::Assume-Yes \"true\";" > /etc/apt/apt.conf.d/90assumeyes
diff --git a/images/linux/scripts/base/repos.sh b/images/linux/scripts/base/repos.sh
index 9005d32aa..a489c62ae 100644
--- a/images/linux/scripts/base/repos.sh
+++ b/images/linux/scripts/base/repos.sh
@@ -16,4 +16,7 @@ curl -L https://packages.microsoft.com/keys/microsoft.asc | apt-key add -
 
 curl https://packages.microsoft.com/keys/microsoft.asc | gpg --dearmor > microsoft.gpg
 mv microsoft.gpg /etc/apt/trusted.gpg.d/microsoft.gpg
-apt-get update
\ No newline at end of file
+
+# update
+apt-get -yq update
+apt-get -yq dist-upgrade
diff --git a/images/linux/ubuntu1604.json b/images/linux/ubuntu1604.json
index a519b30d9..6168f4bf8 100644
--- a/images/linux/ubuntu1604.json
+++ b/images/linux/ubuntu1604.json
@@ -72,6 +72,9 @@
         {
             "type": "shell",
             "script": "{{template_dir}}/scripts/base/apt.sh",
+            "environment_vars": [
+                "DEBIAN_FRONTEND=noninteractive"
+            ],
             "execute_command": "sudo sh -c '{{ .Vars }} {{ .Path }}'"
         },
         {
diff --git a/images/linux/ubuntu1804.json b/images/linux/ubuntu1804.json
index 977464e3a..6fc3f813f 100644
--- a/images/linux/ubuntu1804.json
+++ b/images/linux/ubuntu1804.json
@@ -75,6 +75,9 @@
         {
             "type": "shell",
             "script": "{{template_dir}}/scripts/base/apt.sh",
+            "environment_vars": [
+                "DEBIAN_FRONTEND=noninteractive"
+            ],
             "execute_command": "sudo sh -c '{{ .Vars }} {{ .Path }}'"
         },
         {
diff --git a/images/linux/ubuntu2004.json b/images/linux/ubuntu2004.json
index 54bc6fb57..bde4024c0 100644
--- a/images/linux/ubuntu2004.json
+++ b/images/linux/ubuntu2004.json
@@ -74,9 +74,20 @@
             ],
             "execute_command": "sudo sh -c '{{ .Vars }} {{ .Path }}'"
         },
+        {
+            "type": "shell",
+            "expect_disconnect": true,
+            "scripts": [
+                "{{template_dir}}/scripts/base/reboot.sh"
+            ],
+            "execute_command": "/bin/sh -c '{{ .Vars }} {{ .Path }}'"
+        },
         {
             "type": "shell",
             "script": "{{template_dir}}/scripts/base/apt.sh",
+            "environment_vars": [
+                "DEBIAN_FRONTEND=noninteractive"
+            ],
             "execute_command": "sudo sh -c '{{ .Vars }} {{ .Path }}'"
         },
         {