From 41665e5add665ed819e40c5692793a4393d45bf1 Mon Sep 17 00:00:00 2001
From: Yona Cohen <ycohen@nvidia.com>
Date: Mon, 15 Aug 2022 15:05:41 +0300
Subject: [PATCH 1/8] Implemented secure upgrade

---
 Makefile.work                                 |  3 ++
 build_image.sh                                |  2 +-
 .../build_templates/sonic_debian_extension.j2 |  3 ++
 installer/sharch_body.sh                      |  8 +++-
 onie-mk-demo.sh                               | 41 +++++++++++++++++++
 rules/config                                  |  8 ++++
 scripts/sign_image_dev.sh                     | 14 +++++++
 slave.mk                                      |  9 ++++
 8 files changed, 85 insertions(+), 3 deletions(-)
 create mode 100755 scripts/sign_image_dev.sh

diff --git a/Makefile.work b/Makefile.work
index 3f64531ab11..f7111185ba3 100644
--- a/Makefile.work
+++ b/Makefile.work
@@ -400,6 +400,9 @@ SONIC_BUILD_INSTRUCTION :=  make \
                            SONIC_ENABLE_IMAGE_SIGNATURE=$(ENABLE_IMAGE_SIGNATURE) \
                            SONIC_ENABLE_SECUREBOOT_SIGNATURE=$(SONIC_ENABLE_SECUREBOOT_SIGNATURE) \
                            SONIC_DEFAULT_CONTAINER_REGISTRY=$(DEFAULT_CONTAINER_REGISTRY) \
+                           SECURE_UPGRADE_MODE=$(SECURE_UPGRADE_MODE) \
+                           SECURE_UPGRADE_DEV_SIGNING_KEY=$(SECURE_UPGRADE_DEV_SIGNING_KEY) \
+                           SECURE_UPGRADE_DEV_SIGNING_CERT=$(SECURE_UPGRADE_DEV_SIGNING_CERT) \
                            ENABLE_HOST_SERVICE_ON_START=$(ENABLE_HOST_SERVICE_ON_START) \
                            SLAVE_DIR=$(SLAVE_DIR) \
                            ENABLE_AUTO_TECH_SUPPORT=$(ENABLE_AUTO_TECH_SUPPORT) \
diff --git a/build_image.sh b/build_image.sh
index ddf134e845c..6559b89eda6 100755
--- a/build_image.sh
+++ b/build_image.sh
@@ -86,7 +86,7 @@ generate_onie_installer_image()
     ## Note: Don't leave blank between lines. It is single line command.
     ./onie-mk-demo.sh $CONFIGURED_ARCH $TARGET_MACHINE $TARGET_PLATFORM-$TARGET_MACHINE-$ONIEIMAGE_VERSION \
           installer platform/$TARGET_MACHINE/platform.conf $output_file OS $IMAGE_VERSION $ONIE_IMAGE_PART_SIZE \
-          $ONIE_INSTALLER_PAYLOAD
+          $ONIE_INSTALLER_PAYLOAD $SECURE_UPGRADE_DEV_SIGNING_CERT $SECURE_UPGRADE_DEV_SIGNING_KEY
 }
 
 # Generate asic-specific device list
diff --git a/files/build_templates/sonic_debian_extension.j2 b/files/build_templates/sonic_debian_extension.j2
index 4b7a77b3151..bd8a4d3a915 100644
--- a/files/build_templates/sonic_debian_extension.j2
+++ b/files/build_templates/sonic_debian_extension.j2
@@ -78,6 +78,9 @@ fi
 # Update apt's snapshot of its repos
 sudo LANG=C DEBIAN_FRONTEND=noninteractive chroot $FILESYSTEM_ROOT apt-get update
 
+# Install efitools to support secure upgrade
+sudo LANG=C DEBIAN_FRONTEND=noninteractive chroot $FILESYSTEM_ROOT apt-get -y install efitools
+
 # Apply environtment configuration files
 sudo cp $IMAGE_CONFIGS/environment/environment $FILESYSTEM_ROOT/etc/
 sudo cp $IMAGE_CONFIGS/environment/motd $FILESYSTEM_ROOT/etc/
diff --git a/installer/sharch_body.sh b/installer/sharch_body.sh
index e6289371cd5..d2d46c8b09c 100644
--- a/installer/sharch_body.sh
+++ b/installer/sharch_body.sh
@@ -11,7 +11,9 @@
 ##
 
 echo -n "Verifying image checksum ..."
-sha1=$(sed -e '1,/^exit_marker$/d' "$0" | sha1sum | awk '{ print $1 }')
+payload_image_size=%%PAYLOAD_IMAGE_SIZE%%
+
+sha1=$(sed -e '1,/^exit_marker$/d' "$0" | head -c $payload_image_size | sha1sum | awk '{ print $1 }')
 
 payload_sha1=%%IMAGE_SHA1%%
 
@@ -45,7 +47,9 @@ if [ "$(id -u)" = "0" ] ; then
 fi
 cd $tmp_dir
 echo -n "Preparing image archive ..."
-sed -e '1,/^exit_marker$/d' $archive_path | tar xf - || exit 1
+
+sed -e '1,/^exit_marker$/d' $archive_path | head -c $payload_image_size | tar xf - || clean_up 1
+
 echo " OK."
 cd $cur_wd
 if [ -n "$extract" ] ; then
diff --git a/onie-mk-demo.sh b/onie-mk-demo.sh
index 0905673d42c..5b91ccc817f 100755
--- a/onie-mk-demo.sh
+++ b/onie-mk-demo.sh
@@ -14,6 +14,8 @@ output_file=$6
 demo_type=$7
 image_version=$8
 onie_image_part_size=$9
+cert_file=${11}
+key_file=${12}
 
 shift 9
 
@@ -130,7 +132,46 @@ cp $installer_dir/sharch_body.sh $output_file || {
 # Replace variables in the sharch template
 sed -i -e "s/%%IMAGE_SHA1%%/$sha1/" $output_file
 echo -n "."
+tar_size="$(wc -c < "${sharch}")"
 cat $sharch >> $output_file
+sed -i -e "s|%%PAYLOAD_IMAGE_SIZE%%|${tar_size}|" ${output_file}
+echo "secure upgrade flags: SECURE_UPGRADE_MODE = $SECURE_UPGRADE_MODE, \
+SECURE_UPGRADE_DEV_SIGNING_KEY = $SECURE_UPGRADE_DEV_SIGNING_KEY, SECURE_UPGRADE_DEV_SIGNING_CERT = $SECURE_UPGRADE_DEV_SIGNING_CERT"
+
+if [ "$SECURE_UPGRADE_MODE" = "dev" -o "$SECURE_UPGRADE_MODE" = "prod" ]; then
+    CMS_SIG="${tmp_dir}/signature.sig"
+
+    echo "$0 Creating CMS signature for ${output_file} with  ${key_file}. Output file ${CMS_SIG}"
+    DIR="$(dirname "$0")"
+    
+    scripts_dir="${DIR}/scripts"
+    if [ "$SECURE_UPGRADE_MODE" = "dev" ]; then
+        . ${scripts_dir}/sign_image_dev.sh
+        sign_image_dev ${cert_file} ${key_file} ${output_file} ${CMS_SIG} || { 
+        echo "CMS sign error $?"
+        sudo rm -rf ${CMS_SIG}
+        clean_up 1
+    }
+    else # "$SECURE_UPGRADE_MODE" has to be equal to "prod"
+        . ${scripts_dir}/sign_image_${platform}.sh
+        sign_image_prod ${output_file} ${CMS_SIG} || { 
+        echo "CMS sign error $?"
+        sudo rm -rf ${CMS_SIG}
+        clean_up 1
+    }
+    fi
+    
+    [ -f "$CMS_SIG" ] || {
+         echo "Error: CMS signature not created - exiting without signing"
+         clean_up 1
+    }
+    # append signature to binary
+    cat ${CMS_SIG} >> ${output_file} 
+    sudo rm -rf ${CMS_SIG}
+elif [ "$SECURE_UPGRADE_MODE" -ne "no_sign" ]; then
+    echo "SECURE_UPGRADE_MODE not defined or defined as $SECURE_UPGRADE_MODE - build without signing"
+fi
+
 rm -rf $tmp_dir
 echo " Done."
 
diff --git a/rules/config b/rules/config
index 380b28bc7cd..a407c0921df 100644
--- a/rules/config
+++ b/rules/config
@@ -208,6 +208,14 @@ SONIC_ENABLE_IMAGE_SIGNATURE ?= n
 # The absolute path should be provided.
 SONIC_ENABLE_SECUREBOOT_SIGNATURE ?= n
 
+# folloing flags are used for image secure upgrade verification:
+# SECURE_UPGRADE_DEV_SIGNING_KEY - path to development signing key, used for image signing during build
+# SECURE_UPGRADE_DEV_SIGNING_CERT - path to development signing certificate, used for image signing during build
+# SECURE_UPGRADE_MODE - enum value for secure upgrade mode, valid options are "dev", "prod" and "no_sign"
+#SECURE_UPGRADE_DEV_SIGNING_KEY = 
+#SECURE_UPGRADE_DEV_SIGNING_CERT =  
+SECURE_UPGRADE_MODE = "no_sign"
+
 # PACKAGE_URL_PREFIX - the package url prefix
 PACKAGE_URL_PREFIX ?= https://packages.trafficmanager.net/public/packages
 
diff --git a/scripts/sign_image_dev.sh b/scripts/sign_image_dev.sh
new file mode 100755
index 00000000000..1dbc7fa1bbd
--- /dev/null
+++ b/scripts/sign_image_dev.sh
@@ -0,0 +1,14 @@
+sign_image_dev()
+{
+  cert_file=$1
+  key_file=$2
+  image_to_sign=$3
+  cms_sig_out=$4
+  openssl cms -sign -nosmimecap -signer ${cert_file} -inkey ${key_file} -binary -in $image_to_sign -outform pem -out ${cms_sig_out} || {
+      echo "$?: CMS sign error"
+      sudo rm -rf ${cms_sig_out}
+      exit 1
+  }
+  echo "CMS sign OK"
+  return 0
+}
\ No newline at end of file
diff --git a/slave.mk b/slave.mk
index e1f4a0ef80d..0cd26d11bdf 100644
--- a/slave.mk
+++ b/slave.mk
@@ -348,6 +348,9 @@ $(info "USE_NATIVE_DOCKERD_FOR_BUILD"    : "$(SONIC_CONFIG_USE_NATIVE_DOCKERD_FO
 $(info "SONIC_USE_DOCKER_BUILDKIT"       : "$(SONIC_USE_DOCKER_BUILDKIT)")
 $(info "USERNAME"                        : "$(USERNAME)")
 $(info "PASSWORD"                        : "$(PASSWORD)")
+$(info "SECURE_UPGRADE_MODE"             : "$(SECURE_UPGRADE_MODE)")
+$(info "SECURE_UPGRADE_DEV_SIGNING_KEY"  : "$(SECURE_UPGRADE_DEV_SIGNING_KEY)")
+$(info "SECURE_UPGRADE_DEV_SIGNING_CERT" : "$(SECURE_UPGRADE_DEV_SIGNING_CERT)")
 $(info "ENABLE_DHCP_GRAPH_SERVICE"       : "$(ENABLE_DHCP_GRAPH_SERVICE)")
 $(info "SHUTDOWN_BGP_ON_START"           : "$(SHUTDOWN_BGP_ON_START)")
 $(info "ENABLE_PFCWD_ON_START"           : "$(ENABLE_PFCWD_ON_START)")
@@ -1174,6 +1177,9 @@ $(addprefix $(TARGET_PATH)/, $(SONIC_INSTALLERS)) : $(TARGET_PATH)/% : \
 	export enable_organization_extensions="$(ENABLE_ORGANIZATION_EXTENSIONS)"
 	export enable_dhcp_graph_service="$(ENABLE_DHCP_GRAPH_SERVICE)"
 	export enable_ztp="$(ENABLE_ZTP)"
+	export sonic_su_dev_signing_key="$(SECURE_UPGRADE_DEV_SIGNING_KEY)"
+	export sonic_su_dev_signing_cert="$(SECURE_UPGRADE_DEV_SIGNING_CERT)"
+	export sonic_su_mode="$(SECURE_UPGRADE_MODE)"
 	export include_system_telemetry="$(INCLUDE_SYSTEM_TELEMETRY)"
 	export include_restapi="$(INCLUDE_RESTAPI)"
 	export include_nat="$(INCLUDE_NAT)"
@@ -1373,6 +1379,9 @@ $(addprefix $(TARGET_PATH)/, $(SONIC_INSTALLERS)) : $(TARGET_PATH)/% : \
 		TARGET_MACHINE=$(dep_machine) \
 		IMAGE_TYPE=$($*_IMAGE_TYPE) \
 		SONIC_ENABLE_IMAGE_SIGNATURE="$(SONIC_ENABLE_IMAGE_SIGNATURE)" \
+		SECURE_UPGRADE_MODE="$(SECURE_UPGRADE_MODE)" \
+		SECURE_UPGRADE_DEV_SIGNING_KEY="$(SECURE_UPGRADE_DEV_SIGNING_KEY)" \
+		SECURE_UPGRADE_DEV_SIGNING_CERT="$(SECURE_UPGRADE_DEV_SIGNING_CERT)" \
 		SIGNING_KEY="$(SIGNING_KEY)" \
 		SIGNING_CERT="$(SIGNING_CERT)" \
 		CA_CERT="$(CA_CERT)" \

From 5848e45cdfe63f7f2a76d329a47e0929ddf04c0e Mon Sep 17 00:00:00 2001
From: ycoheNvidia <99744138+ycoheNvidia@users.noreply.github.com>
Date: Thu, 18 Aug 2022 17:28:39 +0300
Subject: [PATCH 2/8] Fixed weird binary related sha1 issue

---
 onie-mk-demo.sh | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/onie-mk-demo.sh b/onie-mk-demo.sh
index 5b91ccc817f..4236e2e9983 100755
--- a/onie-mk-demo.sh
+++ b/onie-mk-demo.sh
@@ -122,7 +122,6 @@ echo -n "."
     echo "Error: $sharch not found"
     clean_up 1
 }
-sha1=$(cat $sharch | sha1sum | awk '{print $1}')
 echo -n "."
 cp $installer_dir/sharch_body.sh $output_file || {
     echo "Error: Problems copying sharch_body.sh"
@@ -130,10 +129,11 @@ cp $installer_dir/sharch_body.sh $output_file || {
 }
 
 # Replace variables in the sharch template
-sed -i -e "s/%%IMAGE_SHA1%%/$sha1/" $output_file
 echo -n "."
 tar_size="$(wc -c < "${sharch}")"
 cat $sharch >> $output_file
+sha1=$(sed -e '1,/^exit_marker$/d' "$output_file" | sha1sum | awk '{ print $1 }')
+sed -i -e "s/%%IMAGE_SHA1%%/$sha1/" $output_file
 sed -i -e "s|%%PAYLOAD_IMAGE_SIZE%%|${tar_size}|" ${output_file}
 echo "secure upgrade flags: SECURE_UPGRADE_MODE = $SECURE_UPGRADE_MODE, \
 SECURE_UPGRADE_DEV_SIGNING_KEY = $SECURE_UPGRADE_DEV_SIGNING_KEY, SECURE_UPGRADE_DEV_SIGNING_CERT = $SECURE_UPGRADE_DEV_SIGNING_CERT"
@@ -168,7 +168,7 @@ if [ "$SECURE_UPGRADE_MODE" = "dev" -o "$SECURE_UPGRADE_MODE" = "prod" ]; then
     # append signature to binary
     cat ${CMS_SIG} >> ${output_file} 
     sudo rm -rf ${CMS_SIG}
-elif [ "$SECURE_UPGRADE_MODE" -ne "no_sign" ]; then
+elif [ "$SECURE_UPGRADE_MODE" != "no_sign" ]; then
     echo "SECURE_UPGRADE_MODE not defined or defined as $SECURE_UPGRADE_MODE - build without signing"
 fi
 

From b3fdcb4de2482f740079d4df863fe510fe8d8661 Mon Sep 17 00:00:00 2001
From: ycoheNvidia <99744138+ycoheNvidia@users.noreply.github.com>
Date: Sun, 21 Aug 2022 15:22:35 +0300
Subject: [PATCH 3/8] Fixed sha1 installer issue

---
 onie-mk-demo.sh | 15 ++++++---------
 1 file changed, 6 insertions(+), 9 deletions(-)

diff --git a/onie-mk-demo.sh b/onie-mk-demo.sh
index 4236e2e9983..4391f823839 100755
--- a/onie-mk-demo.sh
+++ b/onie-mk-demo.sh
@@ -122,6 +122,7 @@ echo -n "."
     echo "Error: $sharch not found"
     clean_up 1
 }
+sha1=$(cat $sharch | sha1sum | awk '{print $1}')
 echo -n "."
 cp $installer_dir/sharch_body.sh $output_file || {
     echo "Error: Problems copying sharch_body.sh"
@@ -129,44 +130,40 @@ cp $installer_dir/sharch_body.sh $output_file || {
 }
 
 # Replace variables in the sharch template
+sed -i -e "s/%%IMAGE_SHA1%%/$sha1/" $output_file
 echo -n "."
 tar_size="$(wc -c < "${sharch}")"
-cat $sharch >> $output_file
-sha1=$(sed -e '1,/^exit_marker$/d' "$output_file" | sha1sum | awk '{ print $1 }')
-sed -i -e "s/%%IMAGE_SHA1%%/$sha1/" $output_file
 sed -i -e "s|%%PAYLOAD_IMAGE_SIZE%%|${tar_size}|" ${output_file}
+cat $sharch >> $output_file
 echo "secure upgrade flags: SECURE_UPGRADE_MODE = $SECURE_UPGRADE_MODE, \
 SECURE_UPGRADE_DEV_SIGNING_KEY = $SECURE_UPGRADE_DEV_SIGNING_KEY, SECURE_UPGRADE_DEV_SIGNING_CERT = $SECURE_UPGRADE_DEV_SIGNING_CERT"
 
 if [ "$SECURE_UPGRADE_MODE" = "dev" -o "$SECURE_UPGRADE_MODE" = "prod" ]; then
     CMS_SIG="${tmp_dir}/signature.sig"
-
     echo "$0 Creating CMS signature for ${output_file} with  ${key_file}. Output file ${CMS_SIG}"
     DIR="$(dirname "$0")"
-    
     scripts_dir="${DIR}/scripts"
     if [ "$SECURE_UPGRADE_MODE" = "dev" ]; then
         . ${scripts_dir}/sign_image_dev.sh
-        sign_image_dev ${cert_file} ${key_file} ${output_file} ${CMS_SIG} || { 
+        sign_image_dev ${cert_file} ${key_file} ${output_file} ${CMS_SIG} || {
         echo "CMS sign error $?"
         sudo rm -rf ${CMS_SIG}
         clean_up 1
     }
     else # "$SECURE_UPGRADE_MODE" has to be equal to "prod"
         . ${scripts_dir}/sign_image_${platform}.sh
-        sign_image_prod ${output_file} ${CMS_SIG} || { 
+        sign_image_prod ${output_file} ${CMS_SIG} || {
         echo "CMS sign error $?"
         sudo rm -rf ${CMS_SIG}
         clean_up 1
     }
     fi
-    
     [ -f "$CMS_SIG" ] || {
          echo "Error: CMS signature not created - exiting without signing"
          clean_up 1
     }
     # append signature to binary
-    cat ${CMS_SIG} >> ${output_file} 
+    cat ${CMS_SIG} >> ${output_file}
     sudo rm -rf ${CMS_SIG}
 elif [ "$SECURE_UPGRADE_MODE" != "no_sign" ]; then
     echo "SECURE_UPGRADE_MODE not defined or defined as $SECURE_UPGRADE_MODE - build without signing"

From e11fb1dc5bc94ebcefbf87af5fd05c480f622bb7 Mon Sep 17 00:00:00 2001
From: ycoheNvidia <99744138+ycoheNvidia@users.noreply.github.com>
Date: Tue, 23 Aug 2022 11:23:53 +0300
Subject: [PATCH 4/8] Fixed minor sharch_body.sh issue

---
 installer/sharch_body.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/installer/sharch_body.sh b/installer/sharch_body.sh
index d2d46c8b09c..9683b4692dc 100644
--- a/installer/sharch_body.sh
+++ b/installer/sharch_body.sh
@@ -48,7 +48,7 @@ fi
 cd $tmp_dir
 echo -n "Preparing image archive ..."
 
-sed -e '1,/^exit_marker$/d' $archive_path | head -c $payload_image_size | tar xf - || clean_up 1
+sed -e '1,/^exit_marker$/d' $archive_path | head -c $payload_image_size | tar xf - || exit 1
 
 echo " OK."
 cd $cur_wd

From 56c9fa68746f9433b6c3128a25af2e56c792a7db Mon Sep 17 00:00:00 2001
From: ycoheNvidia <99744138+ycoheNvidia@users.noreply.github.com>
Date: Tue, 23 Aug 2022 15:09:40 +0300
Subject: [PATCH 5/8] minor typo fix

---
 rules/config | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/rules/config b/rules/config
index a407c0921df..42579983099 100644
--- a/rules/config
+++ b/rules/config
@@ -208,7 +208,7 @@ SONIC_ENABLE_IMAGE_SIGNATURE ?= n
 # The absolute path should be provided.
 SONIC_ENABLE_SECUREBOOT_SIGNATURE ?= n
 
-# folloing flags are used for image secure upgrade verification:
+# following flags are used for image secure upgrade verification:
 # SECURE_UPGRADE_DEV_SIGNING_KEY - path to development signing key, used for image signing during build
 # SECURE_UPGRADE_DEV_SIGNING_CERT - path to development signing certificate, used for image signing during build
 # SECURE_UPGRADE_MODE - enum value for secure upgrade mode, valid options are "dev", "prod" and "no_sign"

From f44a67b9e25179642b9d213cf6972ba03ed431c1 Mon Sep 17 00:00:00 2001
From: ycoheNvidia <99744138+ycoheNvidia@users.noreply.github.com>
Date: Tue, 23 Aug 2022 15:09:59 +0300
Subject: [PATCH 6/8] newline

---
 scripts/sign_image_dev.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/scripts/sign_image_dev.sh b/scripts/sign_image_dev.sh
index 1dbc7fa1bbd..f439243864c 100755
--- a/scripts/sign_image_dev.sh
+++ b/scripts/sign_image_dev.sh
@@ -11,4 +11,4 @@ sign_image_dev()
   }
   echo "CMS sign OK"
   return 0
-}
\ No newline at end of file
+}

From 02d126aafa3be6117adf26306712a70b8d8b6cf8 Mon Sep 17 00:00:00 2001
From: Yona Cohen <ycohen@nvidia.com>
Date: Mon, 9 Jan 2023 16:17:46 +0200
Subject: [PATCH 7/8] fixed secure upgrade image creation bug copying
 irrelevant files to image

---
 onie-mk-demo.sh | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/onie-mk-demo.sh b/onie-mk-demo.sh
index 4391f823839..86464bd021e 100755
--- a/onie-mk-demo.sh
+++ b/onie-mk-demo.sh
@@ -14,6 +14,7 @@ output_file=$6
 demo_type=$7
 image_version=$8
 onie_image_part_size=$9
+onie_installer_payload=${10}
 cert_file=${11}
 key_file=${12}
 
@@ -102,7 +103,7 @@ sed -i -e "s/%%DEMO_TYPE%%/$demo_type/g" \
        -e "s@%%OUTPUT_RAW_IMAGE%%@$output_raw_image@" \
     $tmp_installdir/install.sh || clean_up 1
 echo -n "."
-cp -r $* $tmp_installdir || clean_up 1
+cp -r $onie_installer_payload $tmp_installdir || clean_up 1
 echo -n "."
 [ -r "$platform_conf" ] && {
     cp $platform_conf $tmp_installdir || clean_up 1

From 510c7927ac6db0ff1e87edca0b0a25c98afceb2c Mon Sep 17 00:00:00 2001
From: Yona Cohen <ycohen@nvidia.com>
Date: Tue, 4 Apr 2023 14:06:34 +0300
Subject: [PATCH 8/8] Improved prints and remove image on build failure

---
 onie-mk-demo.sh | 44 +++++++++++++++++++++++++++++---------------
 1 file changed, 29 insertions(+), 15 deletions(-)

diff --git a/onie-mk-demo.sh b/onie-mk-demo.sh
index 86464bd021e..9ce30201e8c 100755
--- a/onie-mk-demo.sh
+++ b/onie-mk-demo.sh
@@ -64,6 +64,13 @@ tmp_dir=
 clean_up()
 {
     rm -rf $tmp_dir
+    if [ -n "$2" ]; then
+        rm -rf "$2"
+        if [ -n "$3" ];then
+            rm -rf "$3"
+        fi
+        echo "Error: CMS signature not created - exiting without signing"
+    fi
     exit $1
 }
 
@@ -134,31 +141,38 @@ cp $installer_dir/sharch_body.sh $output_file || {
 sed -i -e "s/%%IMAGE_SHA1%%/$sha1/" $output_file
 echo -n "."
 tar_size="$(wc -c < "${sharch}")"
-sed -i -e "s|%%PAYLOAD_IMAGE_SIZE%%|${tar_size}|" ${output_file}
 cat $sharch >> $output_file
+sed -i -e "s|%%PAYLOAD_IMAGE_SIZE%%|${tar_size}|" ${output_file}
 echo "secure upgrade flags: SECURE_UPGRADE_MODE = $SECURE_UPGRADE_MODE, \
 SECURE_UPGRADE_DEV_SIGNING_KEY = $SECURE_UPGRADE_DEV_SIGNING_KEY, SECURE_UPGRADE_DEV_SIGNING_CERT = $SECURE_UPGRADE_DEV_SIGNING_CERT"
 
 if [ "$SECURE_UPGRADE_MODE" = "dev" -o "$SECURE_UPGRADE_MODE" = "prod" ]; then
     CMS_SIG="${tmp_dir}/signature.sig"
-    echo "$0 Creating CMS signature for ${output_file} with  ${key_file}. Output file ${CMS_SIG}"
     DIR="$(dirname "$0")"
     scripts_dir="${DIR}/scripts"
+    echo "$0 $SECURE_UPGRADE_MODE signing - creating CMS signature for ${output_file}. Output file ${CMS_SIG}"
+
     if [ "$SECURE_UPGRADE_MODE" = "dev" ]; then
-        . ${scripts_dir}/sign_image_dev.sh
-        sign_image_dev ${cert_file} ${key_file} ${output_file} ${CMS_SIG} || {
-        echo "CMS sign error $?"
-        sudo rm -rf ${CMS_SIG}
-        clean_up 1
-    }
+        echo "$0 dev keyfile location: ${key_file}."
+        . ${scripts_dir}/sign_image_dev.sh || {
+            echo "dev sign script ${scripts_dir}/sign_image_dev.sh not found"
+            clean_up 1 ${output_file}
+        }
+        sign_image_dev ${cert_file} ${key_file} ${output_file} ${CMS_SIG} || { 
+            echo "CMS sign error $?"
+            clean_up 1 ${CMS_SIG} ${output_file}
+        }
     else # "$SECURE_UPGRADE_MODE" has to be equal to "prod"
-        . ${scripts_dir}/sign_image_${platform}.sh
-        sign_image_prod ${output_file} ${CMS_SIG} || {
-        echo "CMS sign error $?"
-        sudo rm -rf ${CMS_SIG}
-        clean_up 1
-    }
+        . ${scripts_dir}/sign_image_${machine}.sh || {
+            echo "prod sign script ${scripts_dir}/sign_image_${machine}.sh not found"
+            clean_up 1 ${output_file}
+        }
+        sign_image_prod ${output_file} ${CMS_SIG} ${SECURE_UPGRADE_MODE} || { 
+            echo "CMS sign error $?"
+            clean_up 1 ${CMS_SIG} ${output_file}
+        }
     fi
+    
     [ -f "$CMS_SIG" ] || {
          echo "Error: CMS signature not created - exiting without signing"
          clean_up 1
@@ -166,7 +180,7 @@ if [ "$SECURE_UPGRADE_MODE" = "dev" -o "$SECURE_UPGRADE_MODE" = "prod" ]; then
     # append signature to binary
     cat ${CMS_SIG} >> ${output_file}
     sudo rm -rf ${CMS_SIG}
-elif [ "$SECURE_UPGRADE_MODE" != "no_sign" ]; then
+elif [ "$SECURE_UPGRADE_MODE" -ne "no_sign" ]; then
     echo "SECURE_UPGRADE_MODE not defined or defined as $SECURE_UPGRADE_MODE - build without signing"
 fi