create-img

#!/bin/bash

#
# SPDX-License-Identifier: Apache-2.0
#
# Name:
#	Create image utility
#
# Authors:
#	2017	Zack YL Shih	<ZackYL.Shih@moxa.com>
#	2018	Fero JD Zhou	<FeroJD.Zhou@moxa.com>
#	2019	Ken CJ Chou	<KenCJ.Chou@moxa.com>
#

set -e

VERSION=1.3.0

usage() {
	echo -e "Usage:"
	echo -e "  # ${0} <image_file> <partition_options1> [<partition_options2> ...]"
	echo -e ""
	echo -e "Parameters:"
	echo -e "  image_file"
	echo -e "    The image filename to be created"
	echo -e ""
	echo -e "  partition_optionsX"
	echo -e "    A partition option is in the format:"
	echo -e "      <part_no>,<part_size>,<boot_flag>,<part_type>,<fs_type>"
	echo -e "    For example:"
	echo -e "      1,32,y,b,vfat"
	echo -e "    this means partition number 1 is a partition with"
	echo -e "    32MB size, boot flag enabled, and vfat filesystem"
	echo -e ""
	echo -e "Example:"
	echo -e "  Create an image named \"output.img\" with 2 partitions"
	echo -e "  whose partition 1 is 32MB size, boot flag enabled, and vfat filesystem"
	echo -e "  and partition 2 is 100MB size, boot flag disabled, and ext4 filesystem"
	echo -e "    # ./create-img output.img \"1,32,y,b,vfat\" \"2,100,n,83,ext4\""
	echo -e ""
	return 0
}

log_msg() {
	local msg=${1}

	echo -ne "\e[1;33m[create-img]\e[0m "
	echo "${msg}"
	return 0
}

is_number() {
	local string=${1}

	case "${string}" in
	""|*[!0-9]*)
		return 1
		;;
	*)
		return 0
		;;
	esac
}

PARTITION_TYPES=" 0 1 2 3 4 5 6 7 8 9 a b c e f 10 11 12 14 16 17 18 "\
"1b 1c 1e 24 27 39 3c 40 41 42 4d 4e 4f 50 51 52 53 54 55 56 5c 61 63 64 65 "\
"70 75 80 81 82 83 84 85 86 87 88 8e 93 94 9f a0 a5 a6 a7 a8 a9 ab af b7 b8 "\
"bb bc be bf c1 c4 c6 c7 da db de df e1 e3 e4 ea eb ee ef f0 f1 f2 f4 fb fc "\
"fd fe ff "

is_valid_partition_type() {
	local part_type=${1}

	if echo "${PARTITION_TYPES}" | grep -q " ${part_type} "; then
		return 0
	else
		return 1
	fi
}

is_valid_fs_type() {
	local fs_type=${1}
	# Currently we only support raw, vfat and ext4
	local fs_support_list="raw vfat ext4"
	local i
	for i in ${fs_support_list}; do
		if [ "${fs_type}" == "${i}" ]; then
			return 0
		fi
	done

	return 1
}

get_value() {
	local tmp
	tmp="$(echo "${1}" | grep "^${2}=")"
	echo "${tmp##*=}"
	return 0
}

get_image_info() {
	local imgfile=${1}
	local fdisk_l_output sec_size num_part part_info
	local line count tmp

	echo "IMAGE_FILE=${imgfile}"
	echo "IMAGE_FILENAME=${imgfile##*/}"

	fdisk_l_output="$(fdisk -l ${imgfile})"

	# get sector size
	sec_size="$(echo "${fdisk_l_output}" | grep "^Units:")"
	sec_size="${sec_size##*= }"
	sec_size="${sec_size%% bytes}"
	echo "IMAGE_SECTOR_SIZE=${sec_size}"

	# get partition information
	part_info="$(echo "${fdisk_l_output}" | grep "^${imgfile}[0-9]")"
	num_part=$(echo "${part_info}" | wc -l)
	echo "IMAGE_NUM_OF_PARTITIONS=${num_part}"

	count=0
	while read -r line; do
		count=$((count + 1))
		# bootflag is set if the line contains '*'
		if echo "${line}" | grep -q "*"; then
			line="$(echo "${line}" | sed 's/*//')"
			echo "IMAGE_P${count}_BOOTFLAG=y"
		else
			echo "IMAGE_P${count}_BOOTFLAG=n"
		fi
		tmp="$(echo "${line}" | awk '{print $1}')"
		echo "IMAGE_P${count}_NUMBER=${tmp##*${imgfile}}"
		tmp="$(echo "${line}" | awk '{print $2}')"
		echo "IMAGE_P${count}_OFFSET=$((tmp * sec_size))"
		tmp="$(echo "${line}" | awk '{print $4}')"
		echo "IMAGE_P${count}_SIZE=$((tmp * sec_size))"
		tmp="$(echo "${line}" | awk '{print $6}')"
		echo "IMAGE_P${count}_FSTYPE=${tmp}"
	done <<< "${part_info}"
	return 0
}

extract_information() {
	local num_of_partitions
	local img_size=0
	local fdisk_cmd=""

	num_of_partitions=$#
	for i in $(seq 1 ${num_of_partitions}); do
		# "<part_no>,<part_size>,<boot_flag>,<part_type>,<fs_type>"
		local part_no="$(echo "${!i}" | awk -F, '{print $1}')"
		local part_size="$(echo "${!i}" | awk -F, '{print $2}')"
		local boot_flag="$(echo "${!i}" | awk -F, '{print $3}')"
		local part_type="$(echo "${!i}" | awk -F, '{print $4}')"
		local fs_type="$(echo "${!i}" | awk -F, '{print $5}')"

		# Check if values are valid
		if ! is_number "${part_no}"; then
			echo -e "Partition number \"${part_no}\" not valid" >&2
			return 1
		fi
		if ! is_number "${part_size}"; then
			echo -e "Partition size \"${part_size}\" not valid" >&2
			return 1
		fi
		if [ "${boot_flag}" != "y" ] && [ "${boot_flag}" != "n" ]; then
			echo -e "Boot flag should be \"y\" or \"n\"" >&2
			return 1
		fi
		if ! is_valid_partition_type "${part_type}"; then
			echo -e "Partition type \"${part_type}\" not valid" >&2
			return 1
		fi
		if ! is_valid_fs_type "${fs_type}"; then
			echo -e "Filesystem type \"${fs_type}\" not valid" >&2
			return 1
		fi

		# Count image total size
		img_size="$(( img_size + part_size ))"

		# Generate command for fdisk
		fdisk_cmd="${fdisk_cmd}n\np\n${part_no}\n\n"
		if [ ${i} -eq 1 ]; then
			fdisk_cmd="${fdisk_cmd}+${part_size}M\n"
			if [ "${boot_flag}" == "y" ]; then
				fdisk_cmd="${fdisk_cmd}a\n"
			fi
			fdisk_cmd="${fdisk_cmd}t\n${part_type}\n"
		elif [ ${i} -eq ${num_of_partitions} ]; then
			fdisk_cmd="${fdisk_cmd}\n"
			if [ "${boot_flag}" == "y" ]; then
				fdisk_cmd="${fdisk_cmd}a\n${part_no}\n"
			fi
			fdisk_cmd="${fdisk_cmd}t\n${part_no}\n${part_type}\n"
		else
			fdisk_cmd="${fdisk_cmd}+${part_size}M\n"
			if [ "${boot_flag}" == "y" ]; then
				fdisk_cmd="${fdisk_cmd}a\n${part_no}\n"
			fi
			fdisk_cmd="${fdisk_cmd}t\n${part_no}\n${part_type}\n"
		fi
	done
	fdisk_cmd="${fdisk_cmd}w\n"

	echo "Num_of_Partitions=${num_of_partitions}"
	echo "Image_Size=${img_size}"
	echo "Fdisk_Command=${fdisk_cmd}"
	return 0
}

format_partitions() {
	local img=${1}
	local img_info i n_part

	shift 1
	img_info="$(get_image_info "${img}")"

	n_part=$(get_value "${img_info}" "IMAGE_NUM_OF_PARTITIONS")
	for i in $(seq 1 ${n_part}); do
		# "<part_no>,<part_size>,<boot_flag>,<part_type>,<fs_type>"
		local fs_type="$(echo "${!i}" | awk -F, '{print $5}')"
		local p_offset=$(get_value "${img_info}" "IMAGE_P${i}_OFFSET")
		local p_size=$(get_value "${img_info}" "IMAGE_P${i}_SIZE")
		local loopdev

		loopdev="$(losetup -f)"
		losetup -o ${p_offset} --sizelimit ${p_size} ${loopdev} ${img}

		# Perform formation for different fs_type
		#   vfat: format
		#   ext4: format with journal
		#   raw: do nothing
		if [ "${fs_type}" == "ext4" ]; then
			mkfs -t ${fs_type} -J size=8 ${loopdev} >/dev/null 2>&1
		elif [ "${fs_type}" == "vfat" ]; then
			mkfs -t ${fs_type} ${loopdev} >/dev/null
		fi

		losetup -d ${loopdev}
	done
	return 0
}

main() {
	local image=${1}
	local extracted_info img_size fdisk_cmd

	if [ "${image}" == "-v" ] || [ "${image}" == "--version" ]; then
		echo "${VERSION}"
		return 0
	fi
	shift 1

	if [ -f ${image} ]; then
		echo "Error: file \"${image}\" already exist" >&2
		return 1
	fi

	log_msg "Extract information from options"
	extracted_info="$(extract_information "$@")"
	img_size="$(get_value "${extracted_info}" "Image_Size")"
	fdisk_cmd="$(get_value "${extracted_info}" "Fdisk_Command")"

	log_msg "Generating image file"
	truncate -s $((img_size * 1024 * 1024)) ${image}

	log_msg "Creating partitions"
	echo -ne "${fdisk_cmd}" | fdisk "${image}" >/dev/null

	log_msg "Formating partitions"
	format_partitions "${image}" "$@"

	log_msg "OK"
	return 0
}

if [ "${EUID}" -ne 0 ]; then
	echo "Please run as root" >&2
	exit 1
fi

if [ $# -lt 2 ]; then
	usage >&2
	exit 1
fi

main "$@"
exit 0