Current File : //bin/usbcopy

#!/usr/bin/bash
#
# Copyright (c) 2010, 2016, Oracle and/or its affiliates. All rights reserved.
#

# Make sure we get the right versions of the commands
PATH=/usr/sbin:/usr/bin

# Ensure that commands whose output we're parsing don't localize numeric
# formats on us
export LC_ALL=C

# Platform (either i386 or sparc)
typeset isa=$(uname -p)

function cleanup {

    {

        # unmounting, and uninstalling the lofi'ed devices and
        # cleanup temporary files.

        for devs in $(mount -p | grep $usbdev | awk '{print $1}'); do
            umount -f $devs
        done

        lofiadm "${img}" && \
            lofiadm -d "${img}"

        if [[ -d "${usbmntpt}" ]]; then
            rmdir "${usbmntpt}"
        fi

        rm -f $fdi
    } > /dev/null 2>&1

}

function error_handler {

	print -u2 "\nError:\n"
	print -u2 -r -- "${progname}: $*"
	cleanup

	exit 1

}

function setup_grub2 {

    # There are pre-built MBR stage 1 and UEFI system images on the USB image.
    # These need to be used instead of trying to construct it ourselves in order
    # to avoid grub2 version mismatches.
    mbrimg="$usbmntpt/boot/mbr.img"
    uefiimg="$usbmntpt/boot/uefi.img"

    # cat the (MBR) usb.img and (ESP) uefi.img to the USB device.
    cat $mbrimg $uefiimg > $dev || \
        error_handler "Failed to write ${mbrimg} and ${uefiimg} to ${dev}"

    # Read the geometry of the physical device
    while read pcyl ncyl acyl bcyl nhead nsect secsiz; do
        devblksize=$secsiz
        devblocks=$(($ncyl * $nhead * $nsect))
    done < <(fdisk -G $dev | grep -v "*")

    # Size the embedded MBR image
    mbrnbytes=$(ls -lL $mbrimg | awk '{print $5}')
    mbrnblocks=$(($mbrnbytes / $devblksize))

    # Size the uefi.img
    uefinbytes=$(ls -lL $uefiimg | awk '{print $5}')
    uefinblocks=$(($uefinbytes / $devblksize))

    # fdisk table should end up looking like this.
    # The first line/partition for the ESP is already embedded in the mbr image
    # that we just cat'd to the disk. So we need to add the solaris partition.
    # id   active bhead bsect bcyl ehead esect ecyl rsect         numsect
    # 239  0      0     0     0    0     0     0    mbrnblocks(a) uefinblocks(b)
    # 191  0      0     0     0    0     0     0    (a)+(b):(c)   diskblocks-(c)
    solstart=$(($mbrnblocks + $uefinblocks))
    solnblocks=$(($devblocks - $solstart))
    addsolpart="191:128:0:0:0:0:0:0:$solstart:$solnblocks"
    fdisk -A $addsolpart $dev

    # fdisk is stupid and overwrites the MBR boot code region with mboot so we
    # have to rewrite the first 440 bytes of the MBR image to restore what it
    # so helpfully overwrote.
    dd if=$mbrimg bs=1 count=440 oseek=0 of=$bdev 2>/dev/null|| \
        error_handler "Failed to restore MBR stage1 image from ${usbimg}"

}


function setup_legacy_grub {

    # Install grub stages to usb. The USB image is already lofi mounted so no need
    # to mount the USB device.
    echo Installing grub to USB device $s0cdev
    installgrub -mf $usbmntpt/boot/grub/stage1 $usbmntpt/boot/grub/stage2 $s0cdev > /dev/null

}

#
# Detect image type and validate it can be deployed on current platform.
# Legacy (non dd-able) x86 image can be deployed on x86 machine only,
# deployment of dd-able images is platform independent (e.g. Sparc dd-able
# image can be deployed on x86 machine).
#
# If image type is unknown or it can't be deployed on current platform,
# inform user and abort.
#
# Return 0 for dd-able image, 1 for legacy (non dd-able) x86 image.
#
function is_image_ddable {

	typeset img="$1"

	#
	# First probe for dd-able image. The check is platform independent.
	# x86 dd-able image begins with MBR ending with magic word '55 aa'.
	# Sparc dd-able image begins with Sparc label identified by magic
	# word 'da be'.
	#
	mbrmagic=$(od -j 510 -N 2 -t x1 "$img" | head -1 | cut -d ' ' -f 2,3)
	lblmagic=$(od -j 508 -N 2 -t x1 "$img" | head -1 | cut -d ' ' -f 2,3)

	if [[ "$mbrmagic" == "55 aa" ]] ; then
		print -u1 "Image type: dd-able x86"

		return 0
	fi
	
	if [[ "$lblmagic" == "da be" ]] ; then
		print -u1 "Image type: dd-able Sparc"

		return 0
	fi

	#
	# Legacy (non dd-able) x86 image is image of ufs root filesystem
	# and can be deployed on x86 only.
	#
	imgtype=$(fstyp "$img" 2>/dev/null)
	if [[ "$imgtype" == "ufs" ]] && [[ "$isa" == "i386" ]] ; then
		print -u1 "Image type: legacy x86"

		return 1
	fi

	print -u2 "Error: Image type not recognized."
	exit 1

}

if [ $(id -u) != "0" ] ; then
	echo "Must be root to run."
	exit 1
fi

if [ $# != 1 ]; then
	echo "Usage: $0 <USB image path>"
	exit 1
fi

img=$1
if [ ! -r $img ] && [ ! -c $img ]; then
	echo "Error: $img does not exist."
	exit 1
fi

# Make sure supplied image is compatible with running platform.
is_image_ddable "$img" && isddable=true || isddable=false

# Image size (in MB)
ibytes=$(ls -lL $img | awk '{print $5}')
isz=$((ibytes >> 20))

#nawk script to output the details of plugged in USB drives
i=0
while read p l s m d; do
	phys[$i]=$p
	log[$i]=$l
	size[$i]=$s
	mult[$i]=$m
	desc[$i]=$d
	((i++))
done < <(rmformat 2>/dev/null | nawk 'BEGIN {
      FS = ":";
      bustype = "USB";
      devtype = "Removable";
}
/Logical Node/ {
	lnode = 1;
	node = $2;
}
/Physical Node/ {
	physdev = $2
}
/Connected Device/ {
	devname = $2
}
/Device Type/ {
	type = $2
}
/Bus/ {
	bus = $2;
}
/Size/ && lnode && bus ~ bustype && type ~ devtype {
	size = $2;
	printf("%s\t%s\t%s\t%s\n", physdev, node, size, devname);
	lnode = 0;
}')

echo Found the following USB devices:
j=0
while [ $j -lt $i ]; do
	echo "$j:	${log[$j]}	${size[$j]} ${mult[$j]}	${desc[$j]}"
	((j++))
done
while read -p "Enter the number of your choice: " choice; do
	if [ -z "${choice}" ]; then
		continue
	fi
	if [ $choice -lt 0 ] || [ $choice -ge $i ]; then
		echo "Invalid choice"
		continue
	fi
	break
done

dev=${log[$choice]}
bdev=${dev/rdsk/dsk}

#
# From character device name representing the whole disk (ends with 'p0' on i386
# and with 's2' on sparc), determine:
#  - slice 0 character device name
#  - character device name representing target which dd(1m) will populate
#    with USB image:
#    - x86  : s0 (for legacy image) or p0 (for dd-able image) 
#    - Sparc: s2
#
if [[ "$isa" == "i386" ]]; then
	s0cdev=${dev/p0/s0}

	if $isddable ; then
		ddcdev="$dev"
	else
		ddcdev="$s0cdev"
	fi

	p1cdev=${dev/p0/p1}
	p1bdev=${p1cdev/rdsk/dsk}
	p2cdev=${dev/p0/p2}
	p2bdev=${p2cdev/rdsk/dsk}
else
	s0cdev=${dev/s2/s0}
	ddcdev="$dev"
fi

s0bdev=${s0cdev/rdsk/dsk}
mountdev=${s0bdev/s0}
if [ ! -b $s0bdev ] || [ ! -c $s0cdev ]; then
	echo "Missing device nodes for $dev"
	exit 1
fi

if [ ! -n "$dev" ]; then
       echo INFORMATION: No USB selected/found.. Please plug in and try again
       exit 1
fi

sz=${size[$choice]}
multiplier=${mult[$choice]}

case "$multiplier" in
    GB) sz=$(( ${sz%%.*} * 1000 + ${sz##*.} * 100 ));;
    MB) sz=${sz%%.*};;
    *)  echo "Unknown capacity indicator: '$multiplier'"; exit 1;;
esac

while true;
do
       echo ""
       echo WARNING: All data on your USB storage will be lost.
       echo Are you sure you want to install to
       echo -n ${desc[$choice]}, $sz MB at $dev ?
       read -p "  (y/n) " yn
       case $yn in
       [yY]* )
	       if [ ! -w $dev ]; then
		   echo "Device $dev is not writable"
		   echo "Installation aborted"
		   exit 1
	       fi
	       if [ $sz -lt $isz ]; then
		   echo "Image ($isz MB) is larger than the device ($sz MB)"
		   echo "Installation aborted"
		   exit 1
	       fi
               break ;;
       [nN]* )
       	       echo "Installation aborted"
	       exit 0 ;;
       * )
       	       echo Invalid choice.. Exiting
	       exit 0 ;;
       esac
done

# Ensure we have things unmounted
for devs in $(mount -p | grep $mountdev | awk '{print $1}');
do    
	umount -f $devs > /dev/null 2>&1                                                                    
done

# Drop in config file to ensure HAL ignores the disk, otherwise it will attempt
# to mount it after the partition table is rewritten
fdi=/etc/hal/fdi/preprobe/10osvendor/99-usbcopy.fdi
trap cleanup SIGTERM EXIT

cat >$fdi << EOF
<?xml version="1.0" encoding="UTF-8"?>
<deviceinfo version="0.2">
  <device>
    <match key="block.device" string="$bdev">
      <merge key="info.ignore" type="bool">true</merge>
    </match>
  </device>
</deviceinfo>
EOF

#
# If deploying on Sparc, make sure that USB stick is labeled with SMI label,
# since it was observed that EFI label (e.g. created by MacOS) prevents
# dd from being able to access the whole device due to the fact that
# slice 2 is regular slice (not the 'backup' one).
# This is not a problem on x86 where USB stick is being accessed via p0
# device node, since p0 always covers the whole available space on the device.
#
if [[ "$isa" == "sparc" ]]; then
	disk_name=${s0cdev#/dev/rdsk/}
	disk_name=${disk_name%s0}

	format -e -L vtoc -d "$disk_name" > /dev/null
	(( $? != 0 )) && { echo Failed to label USB stick. ; exit 1; }
fi

# Legacy x86 image needs more work.
if ! $isddable ; then
	# lofi mount the USB image file into a temp dir
	usbmntpt="$(mktemp -d)"
	usbdev="$(lofiadm -a "${img}")" || \
	    error_handler "Failed to lofiadm ${img}"
	mount -F ufs -o ro $usbdev $usbmntpt || \
	    error_handler "Failed to mount ${usbdev} on ${usbmntpt}" 

	haslegacygrub=false
	if [ -f $usbmntpt/boot/grub/grub.cfg ] ; then
	    # GRUB2
	    setup_grub2
	    zerodev=${p2bdev}
	elif [ -f $usbmntpt/boot/grub/menu.lst ] ; then
	    # Legacy GRUB
	    # Install fdisk table with Solaris using entire disk, default VTOC
	    haslegacygrub=true
	    fdisk -B $dev
	    zerodev=${p1bdev}
	else
	    error_handler "Failed to find boot loader configuration"
	fi

	# Remove the PBR from the Solaris partition to ensure no stray data confuses
	# some BIOSes
	dd if=/dev/zero of=${zerodev} count=1 > /dev/null 2>&1
	(( $? != 0 )) && { echo Failed to zero Solaris PBR on USB device. ; exit 1; }
 
	# Now create root partition. We want to find number of cylinders
	# in backup partition from label created by fdisk and then generate
	# root partition using whole disk minus cylinder 1.
	acyls=$(prtvtoc $dev | awk '/accessible/{print $2}')
	cyls=$((acyls - 1))
	format -e $dev >/dev/null <<EOF
partition
0
root
wm
1
${cyls}c
label
0
y
EOF
fi

# Copy image to USB.
echo "Copying and verifying image to USB device"
SECONDS=0
# For now, copy in 16k chunks, 4MB at a time
# NOTE: To avoid problems, cmb should be >= 1 and bs should evenly
#       divide into csz.
bs=$((16 << 10))
cmb=4

csz=$((cmb << 20))
bcnt=$((csz / bs))
icnt=$((isz / cmb))

maxt=3
i=0
retries=0

while ((i <= icnt)); do
    tries=0
    pcnt="($((i * 100 / icnt))%)"
    s=$((i * bcnt))

    echo -ne "R $i / $icnt  $pcnt \\r"
    imgsum=$(dd bs=$bs count=$bcnt conv=sync \
	iseek=$s if=$img 2>/dev/null | digest -a md5)
    (($? != 0)) && { echo Read from image failed. ; exit 1; }

    while ((tries++ < maxt)); do
	echo -ne "W $i / $icnt  $pcnt \\r"
	recs=$(dd bs=$bs count=$bcnt conv=sync iseek=$s oseek=$s \
	    if=$img of=$ddcdev 2>&1 | awk -F+ '/records out/{print $1}')
	(($? != 0)) && { echo Write to device failed. ; exit 1; }

	echo -ne "V $i / $icnt  $pcnt \\r"

	if [[ "$isa" == "i386" ]]; then
		devsum=$(dd bs=$bs count=$recs iseek=$s if=$ddcdev 2>/dev/null | digest -a md5)
	fi

	if [[ "$isa" == "sparc" ]]; then
		devsum=$(dd bs=$bs count=$recs conv=sync iseek=$s if=$ddcdev 2>/dev/null | digest -a md5)
	fi
	(($? != 0)) && { echo Read from device failed. ; exit 1; }

	[[ "$imgsum" == "$devsum" ]] && break;
    done
    ((tries > 1)) && ((retries++))

    if ((tries > maxt)); then
	echo "Verification failed after $maxt attempts on block $i"
	exit 1
    fi

    ((i++))
done

speed="$((isz / SECONDS)).$((isz * 10 / SECONDS % 10))MB/s"
echo "Finished $isz MB in $SECONDS seconds ($speed)"
echo "$retries block(s) re-written due to verification failure"

# If it is legacy x86 image with legacy GRUB, install bootloader now.
if ! $isddable && $haslegacygrub; then
	setup_legacy_grub
fi

# Cleanup things.
cleanup

echo "Completed copy to USB"
exit 0