Script to re-size stateful partition on ChromeOS

resize-stateful-partition.sh

#!/bin/bash

red="\E[1;31m";
green="\E[1;32m";
yellow="\E[1;33m";
plain="\e[0m";

error() {
    echo -e "${red}${1}${plain}";
}

info() {
    echo -e "${plain}${1}${plain}";
}

warn() {
    echo -e "${yellow}${1}${plain}";
}

abort() {
    [[ -z "$1" ]] && msg="Aborting." || msg="$1";
    error "\n$msg";
    exit 1;
}

confirm() {
    msg="\nPress [Enter] to continue or [Ctrl+C] to abort.";
    [ -n "$1" ] && msg="$1";
    echo -ne "$msg";
    read -u 1;
    echo;
}

insession() {
    p=$$; 
    while [ $p -gt 0 ]; do 
        i=($(ps h -o ppid,cmd --pid $p)); 
        [ "${i[1]}" == "session_manager" ] && return 0;
        p="${i[0]}";
    done;
    return 1; 
}

getmounts() {
    # Get root mountpoints
    [ -z "$1" ] && return 1;
    rootmounts=($(cat /proc/mounts | egrep "^$1" | awk '{print $2}'));
    (
        for m in "${rootmounts[@]}"; do
            [ "$m" == "/" ] && continue;
            # get mapper mounts
            cat /proc/mounts | awk '{print $2}' | egrep "^$m";
        done
    ) | sort | uniq;
}

unmountdev() {
    echo "Attempting to unmount $1";
    local err=0
    for dir in $(getmounts $1 | sort -r); do
        # skip encstateful mount
        src="$(df --output=source | tail -n 1)";
        [ "$src" == "/dev/mapper/encstateful" ] && continue;
        echo "  ${dir}..."
        umount -f "$dir" 2> /dev/null || err=1;
    done
    return $err;
}

svcstarted() {
    status="$(initctl status $1 | awk '{print $2}')"
    [ "$status" == "start/running," ] || return 1;
}

trap abort INT;

[[ $(id -u) == "0" ]] || abort "This script must be run as root.";

[ "$(pwd)" != "/" ] && abort "You must change to the / directory before running this script.";

insession && abort "This script must be run from a VT. Press [Ctrl + Alt + =>] to switch to VT2"

if svcstarted "ui"; then
    msg="The UI service is currently running. We'll need to stop the UI \n";
    msg+="service and re-run this script. Once you stop the UI service, \n";
    msg+="you'll be logged out and will need to log in again.";
    warn "$msg";
    confirm "Press [Enter] to stop the UI service.";
    info "Stopping UI...";
    stop ui || abort;
    exit;
fi


#--| Check for existing partition data |--------------------------------------

rootdev="$(rootdev -d -s)" || abort;
if [ -z "$rootdev" ]; then
    echo "Cannot find root device."
    exit 1
fi

rootcpart="$(cgpt find -n -l ROOT-C "$rootdev")"
if [ "$(cgpt show -i $rootcpart -s "$rootdev")" -gt 1 ]; then
    abort "ROOT-C is not empty."
fi

statepart="$(cgpt find -n -l STATE "$rootdev")"
statedev="$(cgpt find -l STATE "$rootdev")"


warn "************************** WARNING ***************************";
warn "This process will wipe all data in your ChromeOS installation.";
warn "**************************************************************";
echo


#--| Get new partition size from user |---------------------------------------
state_size="$(cgpt show -i ${statepart} -n -s -q ${rootdev})" || abort;
max_partition_size=$(($state_size/1024/1024/2));
rec_partition_size=$(($max_partition_size - 1));

info "Press [Ctrl+C] to abort.\n"
info "Enter the size in gigabytes you want to reserve for your Linux OS. ";
info "Acceptable range: 5-${max_partition_size}"
info "Recommended max:  ${rec_partition_size}\n";

while :; do
    read -u 1 -p "Enter partition size (GB): " partition_size;

    if [[ ! $partition_size =~ ^[0-9]+$ ]]; then
        error "\n\nNumbers only please...\n\n"
        continue
    fi

    if [ $partition_size -lt 5 -o $partition_size -gt $max_partition_size ]; then
        error "\n\nThat number is out of range. Enter a number 5 through $max_partition_size\n\n"
        continue
    fi

    break;
done

while :; do
    read -u 1 -p "Enter label name for new partition: " new_label;

    new_label=$(echo "$new_label" | tr '[:lower:]' '[:upper:]');

    if cgpt find -n -l ${new_label} &> /dev/null; then
        error "Partition ${new_label} already exists."
        continue
    fi

    break;
done


#--| Make calculations |------------------------------------------------------

partition_sector_size=$((partition_size*1024*1024*2));
statestart="$(cgpt show -i "$statepart" -b "$rootdev")"
statesize="$(cgpt show -i "$statepart" -s "$rootdev")"
newstatesize=$((statesize-partition_sector_size))
partition_start=$((statestart+newstatesize))

info "\nNew sizes:\n"
[ ${#new_label} -gt 5 ] && len=$(( ${#new_label} + 1)) || len=6;
printf "${green}%-${len}s${plain} %5d MB\n" "STATE:" $((newstatesize/(2*1024)))
printf "${green}%-${len}s${plain} %5d MB\n\n" "${new_label}:" $((partition_sector_size/(2*1024)))
confirm;


#--| Log us out |----------------------------------------------------------
cd /

info "Stopping services...";
stopservices=(syslog powerd tlsdated tcsd chapsd update-engine warn-collector)
for svc in "${stopservices[@]}"; do
    if svcstarted $svc; then
        sudo stop $svc || abort;
        sleep 1;
    fi
done


#--| Unmount the stateful partition |-----------------------------------------
mounts_list=$(mktemp);
getmounts /dev/mapper/encstateful > "$mounts_list";
getmounts "$statedev" >> "$mounts_list";

info "Unmounting stateful partition...";

# remove all loop devices
if [ -n "`losetup -a`" ]; then
    losetup -D 2> /dev/null || abort;
fi

ok=0;
tries=0;
while :; do
    [ $tries -gt 10 ] && abort "Could not unmount stateful partition"
    ok=1;
    unmountdev /dev/mapper/encstateful || ok=0;
    unmountdev "$statedev" || ok=0;
    dmsetup remove encstateful 2> /dev/null || true;
    [ "$ok" == 1 ] && break;
    tries=$(( tries + 1 ));
done

#--| Resize the stateful partition |------------------------------------------
lastpart="$(sudo cgpt show -n -q /dev/sda | awk '{print $3}' | sort -n | tail -n 1)" || abort;
newpart=$((lastpart + 1));
cmds=("e2fsck -f -p $statedev"
      "resize2fs $statedev ${newstatesize}s"
      "cgpt add -i $statepart -b $statestart -s $newstatesize -l STATE $rootdev"
      "cgpt add -i $newpart -t data -b $partition_start -s $partition_sector_size -l $new_label $rootdev")

info "The following operations will be performed:\n"
for cmd in "${cmds[@]}"; do
    echo -e "${green}>>${plain} ${cmd}";
done
confirm;
info "Resizing stateful partition...";
for cmd in "${cmds[@]}"; do
    $cmd || abort "Failed command:\n  ${plain}$cmd";
done

info "Success! Press [Enter] to reboot...";
read -u 1;

reboot;
exit 0;