1#!/bin/bash
2
3# Copyright 2019 Google Inc. All rights reserved.
4
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8
9#     http://www.apache.org/licenses/LICENSE-2.0
10
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17usage() {
18	echo "USAGE: $0 [flags] image"
19	echo "flags:"
20	echo "  -e,--expand:  expand filesystem to fill device (default: false)"
21	echo "  -h,--help:    show this help (default: false)"
22}
23
24main()
25{
26	if [ ! -e "${image}" ]; then
27		echo "error: can't find image. aborting..."
28		exit 1
29	fi
30
31	init_devs=`lsblk --nodeps -oNAME -n`
32	echo "Reinsert device (to write to) into PC"
33	while true; do
34		devs=`lsblk --nodeps -oNAME -n`
35		new_devs="$(echo -e "${init_devs}\n${devs}" | sort | uniq -u | awk 'NF')"
36		num_devs=`echo "${new_devs}" | wc -l`
37		if [[ "${new_devs}" == "" ]]; then
38			num_devs=0
39		fi
40		if [[ ${num_devs} -gt 1 ]]; then
41			echo "error: too many new devices detected! aborting..."
42			exit 1
43		fi
44		if [[ ${num_devs} -eq 1 ]]; then
45			break
46		fi
47	done
48	blk_dev=${new_devs}
49	# don't inform user we found the block device yet (it's confusing)
50
51	init_devs=${devs}
52	echo "${init_devs}" | grep "${blk_dev}" >/dev/null
53	if [[ $? -ne 0 ]]; then
54		while true; do
55			devs=`lsblk --nodeps -oNAME -n`
56			new_devs="$(echo -e "${init_devs}\n${devs}" | sort | uniq -u | awk 'NF')"
57			num_devs=`echo "${new_devs}" | wc -l`
58			if [[ "${new_devs}" == "" ]]; then
59				num_devs=0
60			fi
61			if [[ ${num_devs} -gt 1 ]]; then
62				echo "error: too many new devices detected! aborting..."
63				exit 1
64			fi
65			if [[ ${num_devs} -eq 1 ]]; then
66				if [[ "${new_devs}" != "${blk_dev}" ]]; then
67					echo "error: block device name mismatch ${new_devs} != ${blk_dev}"
68					echo "Reinsert device (to write to) into PC"
69					blk_dev=${new_devs}
70					new_devs=""
71					continue
72				fi
73				break
74			fi
75		done
76	fi
77	# now inform the user
78	echo "Detected device at /dev/${blk_dev}"
79
80	imgsize=`ls -lah ${image} | awk -F " " {'print $5'}`
81	echo "Ready to write ${imgsize} image to block device at /dev/${blk_dev}..."
82	sudo chmod 666 /dev/${blk_dev}
83	type pv > /dev/null 2>&1
84	if [ $? == 0 ]; then
85		pv ${image} > /dev/${blk_dev}
86	else
87		dd if=${image} of=/dev/${blk_dev} bs=1M conv=sync,noerror status=progress
88	fi
89	if [ $? != 0 ]; then
90		echo "error: failed to write to device. aborting..."
91		exit 1
92	fi
93
94	if [ ${expand} -eq 1 ]; then
95		echo "Expanding partition and filesystem..."
96		part_type=`sudo gdisk -l /dev/${blk_dev}  2>/dev/null | grep ": present" | sed 's/ *\([^:]*\):.*/\1/'`
97		if [ "$part_type" == "MBR" ]; then
98			sudo parted -s /dev/${blk_dev} resizepart 1 100%
99			sudo e2fsck -y -f /dev/${blk_dev}1 >/dev/null 2>&1
100			sudo resize2fs /dev/${blk_dev}1 >/dev/null 2>&1
101		elif [ "$part_type" == "GPT" ]; then
102			parts=`sudo gdisk -l  /dev/${blk_dev} | grep "^Number" -A999 | tail -n +2 | wc -l`
103			FIRST_SECTOR=`sudo gdisk -l /dev/${blk_dev} 2>/dev/null | tail -1 | tr -s ' ' | cut -d" " -f3`
104			sudo sgdisk -d${parts} /dev/${blk_dev} >/dev/null 2>&1
105			sudo sgdisk -a1 -n:${parts}:${FIRST_SECTOR}:- -A:${parts}:set:2 -t:${parts}:8305 -c:${parts}:rootfs /dev/${blk_dev} >/dev/null 2>&1
106			sudo e2fsck -fy /dev/${blk_dev}${parts} >/dev/null 2>&1
107			sudo resize2fs /dev/${blk_dev}${parts} >/dev/null 2>&1
108		fi
109	fi
110	sudo sync /dev/${blk_dev}
111	sudo eject /dev/${blk_dev}
112}
113
114expand=0
115if [ "$1" == "-e" ] || [ "$1" == "--expand" ]; then
116	expand=1
117	shift
118fi
119
120if [ "$1" == "-h" ] || [ "$1" == "--help" ]; then
121	usage
122	exit 0
123fi
124
125image=$1
126if [ "${image}" == "" ]; then
127	usage
128	exit 1
129fi
130
131main "$@"
132