From 17fc7af9a8fe7eb4778238e1e6c831b773614209 Mon Sep 17 00:00:00 2001
From: bin456789 <bin456789@gmail.com>
Date: Sun, 4 Jun 2023 22:10:16 +0800
Subject: [PATCH] =?UTF-8?q?=E6=8F=90=E5=89=8D=E6=B5=8B=E8=AF=95iso/dd?=
 =?UTF-8?q?=E9=95=9C=E5=83=8F=E6=98=AF=E5=90=A6=E6=9C=89=E6=95=88?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 reinstall.sh | 35 ++++++++++++++++++++++++++++++++++-
 trans.sh     |  3 +--
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/reinstall.sh b/reinstall.sh
index e7cc954..1aeac50 100644
--- a/reinstall.sh
+++ b/reinstall.sh
@@ -7,6 +7,11 @@ usage_and_exit() {
     exit 1
 }
 
+error_and_exit() {
+    echo "Error: $1"
+    exit 1
+}
+
 is_in_china() {
     if [ -z $_is_in_china ]; then
         # https://geoip.ubuntu.com/lookup
@@ -16,6 +21,29 @@ is_in_china() {
     return $_is_in_china
 }
 
+test_url() {
+    url=$1
+    expect_type=$2
+    tmp_file=/tmp/reinstall-img-test
+    install_pkg file
+
+    http_code=$(curl -Ls -r 0-1048576 -w "%{http_code}" -o $tmp_file $url)
+    if [ "$http_code" != 200 ] && [ "$http_code" != 206 ]; then
+        error_and_exit "$url not accessible"
+    fi
+
+    mime=$(file -ib $tmp_file | cut -d';' -f1 | cut -d'/' -f2)
+    case $mime in
+    gzip) img_type=gz ;;
+    x-xz) img_type=xz ;;
+    x-iso9660-image) img_type=iso ;;
+    esac
+
+    if ! echo $expect_type | grep -wo "$img_type"; then
+        error_and_exit "$url $mime not supported"
+    fi
+}
+
 setos() {
     local step=$1
     local distro=$2
@@ -83,7 +111,9 @@ setos() {
         fi
 
         filename=$(curl $mirror | grep -oP "ubuntu-$releasever.*?-live-server-$basearch_alt.iso" | head -1)
-        eval ${step}_iso=$mirror$filename
+        iso=$mirror$filename
+        test_url $iso iso
+        eval ${step}_iso=$iso
         eval ${step}_ks=$confhome/user-data
     }
 
@@ -92,6 +122,7 @@ setos() {
             echo "Install Windows need --iso --image-name"
             exit 1
         fi
+        test_url $iso iso
         eval "${step}_iso='$iso'"
         eval "${step}_image_name='$image_name'"
     }
@@ -101,7 +132,9 @@ setos() {
             echo "dd need --ddimg"
             exit 1
         fi
+        test_url $ddimg 'xz|gz'
         eval "${step}_ddimg='$ddimg'"
+        eval "${step}_ddimg_type='$img_type'"
     }
 
     setos_redhat() {
diff --git a/trans.sh b/trans.sh
index ef5af74..ac0da76 100644
--- a/trans.sh
+++ b/trans.sh
@@ -91,8 +91,7 @@ if [ "$distro" = "alpine" ]; then
     exec reboot
 
 elif [ "$distro" = "dd" ]; then
-    filetype=$(echo $ddimg | awk -F. '{print $NF}')
-    case "$filetype" in
+    case "$ddimg_type" in
     gz) prog=gzip ;;
     xz) prog=xz ;;
     esac