bfpart

#!/bin/sh -e

# Copyright (c) 2017, Mellanox Technologies
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# The views and conclusions contained in the software and documentation are those
# of the authors and should not be interpreted as representing official policies,
# either expressed or implied, of the FreeBSD Project.

usage ()
{
    echo "syntax: bfpart [-b MB] [-p MB] [-o]"
}

# Note: These variables are eval'd further down.
# shellcheck disable=SC2034
default_cfg()
{
  DISK0_NAME=/dev/mmcblk0
  DISK0_PART1_SIZE=200
  DISK0_PART1_TYPE=EFI
  DISK0_PART2_SIZE=0
  DISK0_PART2_TYPE=Linux
  DISK0_PART8_PERSIST=KEEP
  DISK0_PART8_SIZE=250
  DISK0_PART8_TYPE=Linux
}

bfcfg_scan()
{
  sector_start_reserve=2048
  sector_end_reserve=33

  # variables in use:
  # disk_name part_size part_type part_persist part_mount disk_id part_uuid
  # disk_sectors sector sector_start total_sectors auto_part cur_persist_uuid
  #
  # local would be used here, but local is not POSIX-compliant.

  # Support up to 16 disks and 16 partitions on each disk.
  for disk in $(seq -s" " 0 15); do
    eval "disk_name=\${DISK${disk}_NAME}"
    # shellcheck disable=SC2154
    [ -z "${disk_name}" ] && continue

    disk_sectors=$(fdisk -l "${disk_name}" | grep "Disk ${disk_name}:" | awk '{print $5}')
    disk_sectors=$((disk_sectors / 512))
    sector_start=${sector_start_reserve}
    total_sectors=$((sector_start_reserve + sector_end_reserve))
    auto_part=255

    eval "DISK${disk}_ID=$(uuidgen)"
    eval "DISK${disk}_START=${sector_start}"

    for part in $(seq -s" " 0 15); do
      eval "part_size=\${DISK${disk}_PART${part}_SIZE}"
      eval "part_type=\${DISK${disk}_PART${part}_TYPE}"
      [ -z "${part_size}" ] || [ -z "${part_type}" ] && continue
      part_type=$(echo "$part_type" | tr '[:lower:]' '[:upper:]')
      eval "part_persist=\${DISK${disk}_PART${part}_PERSIST}"
      eval "part_mount=\${DISK${disk}_PART${part}_MOUNT}"

      # Evaluate and assign partition UUID information.
      if [ "${part_type}" = "EFI" ] || [ "${part_type}" = "C12A7328-F81F-11D2-BA4B-00A0C93EC93B" ]; then
        part_type=C12A7328-F81F-11D2-BA4B-00A0C93EC93B
        part_uuid=3DCADB7E-BCCC-4897-A766-3C070EDD7C25
      else
        if [ "${part_type}" = "LINUX" ] || [ "${part_type}" = "0FC63DAF-8483-4772-8E79-3D69D8477DE4" ]; then
          part_type=0FC63DAF-8483-4772-8E79-3D69D8477DE4
        fi
        if [ -n "${part_persist}" ]; then
          # Convert to upper case.
          part_persist=$(echo "$part_persist" | tr '[:lower:]' '[:upper:]')
          part_uuid=CAFDC47B-3F21-4A78-81CC-F91944081165
        else
          part_uuid=$(uuidgen)
        fi
      fi

      sector_num=$((part_size * 1024 * 1024 / 512))

      cur_persist_uuid="$(sfdisk --part-uuid "${disk_name}" "${part}" || true)"
      if [ ."${part_persist}" = ."KEEP" ]; then
        if [ ."${cur_persist_uuid}" = ."${part_uuid}" ]; then
          sector_start="$(sfdisk -l "${disk_name}" | grep "${disk_name}p${part}" | awk '{print $2}')"
          sector_num="$(sfdisk -l "${disk_name}" | grep "${disk_name}p${part}" | awk '{print $4}')"
        else
          eval "DISK${disk}_PART${part}_PERSIST=CREATE"
        fi
      fi

      if [ "${sector_num}" = "0" ]; then
        auto_part=${part}
      fi

      eval "DISK${disk}_PART${part}_TYPE=${part_type}"
      eval "DISK${disk}_PART${part}_UUID=${part_uuid}"
      eval "DISK${disk}_PART${part}_SECTOR_START=${sector_start}"
      eval "DISK${disk}_PART${part}_SECTOR_NUM=${sector_num}"

      total_sectors=$((total_sectors + sector_num))
      sector_start=$((sector_start + sector_num))
      eval "DISK${disk}_END=$((sector_start - 1))"
    done

    if [ ${total_sectors} -gt ${disk_sectors} ]; then
      echo "total configured size exceeds disk size" 1>&2
      exit 1
    fi

    if [ "$auto_part" -ne 255 ]; then
      part_size=$(((disk_sectors - total_sectors) * 512 / 1024 / 1024))
      eval "DISK${disk}_PART${auto_part}_SIZE=${part_size}"
    fi
  done
}

sfdisk_apply()
{
  # variables in use:
  # sector_start sector_end sector_num part_type part_uuid disk_id persist
  #
  # local would be used here, but local is not POSIX-compliant.

  SFDISK_TMP_CFG=/tmp/disk.sfdisk

  for disk in $(seq -s" " 0 15); do
    rm -f ${SFDISK_TMP_CFG}

    eval "disk_name=\${DISK${disk}_NAME}"
    [ -z "${disk_name}" ] && continue
    eval "disk_id=\${DISK${disk}_ID}"
    eval "sector_start=\${DISK${disk}_START}"
    eval "sector_end=\${DISK${disk}_END}"

    {
      echo "label: gpt"
      # shellcheck disable=SC2154
      echo "label-id: ${disk_id}"
      echo "device: ${disk_name}"
      echo "unit: sectors"
      echo "first-lba: ${sector_start}"
      # shellcheck disable=SC2154
      echo "last-lba: ${sector_end}"
    } >> ${SFDISK_TMP_CFG}

    for part in $(seq -s" " 0 15); do
      eval "part_type=\${DISK${disk}_PART${part}_TYPE}"
      eval "part_uuid=\${DISK${disk}_PART${part}_UUID}"
      eval "sector_start=\${DISK${disk}_PART${part}_SECTOR_START}"
      eval "sector_num=\${DISK${disk}_PART${part}_SECTOR_NUM}"
      eval "persist=\${DISK${disk}_PART${part}_PERSIST}"

      if [ -n "${persist}" ]; then
        persist=$(echo "$persist" | tr '[:upper:]' '[:lower:]')
        echo "PERSIST: ${persist}"
      fi

      [ -z "${part_uuid}" ] && continue

      echo "gpt.dat${part} : start=${sector_start}, size=${sector_num}, type=${part_type}, uuid=${part_uuid}" >> ${SFDISK_TMP_CFG}
    done

    dd if=/dev/zero of="${disk_name}" bs=512 count=1
    sfdisk "${disk_name}" < /tmp/disk.sfdisk
  done
}

BF_CFG_FILE=/etc/bf.cfg

while getopts ":b:p:o" opt; do
    case $opt in
        b)
            boot_size=$OPTARG
            ;;
        p)
            persistent_size=$OPTARG
            ;;
        o)
            # overwrite persistent partition
            override_persist=1
            ;;
        \?)
            usage >&2
            exit 1
        ;;
    esac
done

# Use default config if arguments are provided or config file doesn't exist.
if [ "$#" -ne 0 ] || [ ! -e "${BF_CFG_FILE}" ]; then
  default_cfg
  # shellcheck disable=SC2034
  [ -n "${boot_size}" ] && DISK0_PART1_SIZE=${boot_size}
  # shellcheck disable=SC2034
  [ -n "${persistent_size}" ] && DISK0_PART8_SIZE=${persistent_size}
  # shellcheck disable=SC2034
  [ ."${override_persist}" = ."1" ] && DISK0_PART8_PERSIST=CREATE
else
  # shellcheck source=/dev/null
  . ${BF_CFG_FILE}
  [ -z "${DISK0_NAME}" ] && default_cfg
fi

# First-time scan to figure out all the sizes.
bfcfg_scan

# Second-time scan to adjust the offsets.
bfcfg_scan

# Apply the configuration
sfdisk_apply