#!/usr/bin/env bash
#
# samsung-heimdall-prepare.sh
#
# Prepares a Samsung Odin firmware (AP/BL/CP/CSC *.tar.md5) for flashing with
# Heimdall on Linux:
#   1. verifies the MD5 appended to each .tar.md5
#   2. extracts the images into ./extracted (decompressing *.lz4)
#   3. reads the partition table (PIT) shipped in the CSC archive
#   4. prints the heimdall command that flashes each extracted file to the
#      partition whose "Flash Filename" matches it, the factory USERDATA image
#      last (unless --keep-data is used); run it with the phone in Download mode
#
# It does NOT flash anything and does not talk to the phone. Review the
# printed command, then run it yourself.
#
# Usage (in the firmware folder):
#   ./samsung-heimdall-prepare.sh [--keep-data]
#
#   --keep-data   use HOME_CSC (keeps user data) instead of CSC (factory reset)
#
# Requires: heimdall, tar, md5sum, lz4
#
set -euo pipefail

CSC_PREFIX="CSC_"
[ "${1:-}" = "--keep-data" ] && CSC_PREFIX="HOME_CSC_"

for tool in heimdall tar md5sum lz4; do
    command -v "$tool" >/dev/null 2>&1 || { echo "Error: $tool is not installed." >&2; exit 1; }
done

shopt -s nullglob
files=( AP_*.tar.md5 BL_*.tar.md5 CP_*.tar.md5 ${CSC_PREFIX}*.tar.md5 )
[ "${#files[@]}" -gt 0 ] || { echo "Error: no AP/BL/CP/CSC .tar.md5 files here." >&2; exit 1; }
# Only the CSC archive ships the PIT, HOME_CSC does not.
pit_archives=( CSC_*.tar.md5 )
[ "${#pit_archives[@]}" -eq 1 ] || { echo "Error: expected one CSC_*.tar.md5 file (it holds the PIT)." >&2; exit 1; }
verify=( "${files[@]}" )
[ "$CSC_PREFIX" = "CSC_" ] || verify+=( "${pit_archives[0]}" )

echo "== 1. Verifying MD5 =="
for f in "${verify[@]}"; do
    # A .tar.md5 is a tar archive directly followed by the line
    # "<32-hex md5>  <name>.tar\n" (the output of md5sum appended to the tar).
    tarname="${f%.md5}"
    trailer_len=$(( 32 + 2 + ${#tarname} + 1 ))
    expected="$(tail -c "$trailer_len" "$f" | head -c 32)"
    actual="$(head -c -"$trailer_len" "$f" | md5sum | awk '{print $1}')"
    if [ "$expected" = "$actual" ]; then
        echo "OK   $f"
    else
        echo "FAIL $f (expected $expected, got $actual)" >&2
        exit 1
    fi
done

echo "== 2. Extracting =="
# Start from a clean folder so files left by an earlier run are never reused.
rm -rf extracted
mkdir -p extracted
for f in "${files[@]}"; do
    tar -xf "$f" -C extracted
done
[ "$CSC_PREFIX" = "CSC_" ] || tar -xf "${pit_archives[0]}" -C extracted --wildcards '*.pit'
# Some packages (e.g. CP) store their files read-only; make them writable so
# the .lz4 images can be decompressed next to them.
chmod -R u+w extracted
(
    cd extracted
    for z in *.lz4; do
        rm -f "${z%.lz4}"
        if ! lz4 -d -q "$z" "${z%.lz4}"; then
            echo "Error: could not decompress $z" >&2
            exit 1
        fi
        rm -f "$z"
    done
) || exit 1
ls -l extracted

echo "== 3. Reading the PIT from the firmware =="
# Reading the PIT from the phone would either reboot it or leave a Download
# mode session open that only "--resume" can continue, so read the firmware's
# copy instead: the flash command is then the only one that talks to the phone.
pits=( extracted/*.pit )
[ "${#pits[@]}" -eq 1 ] || { echo "Error: expected one .pit file in the CSC, found ${#pits[@]}." >&2; exit 1; }
if ! heimdall print-pit --file "${pits[0]}" > pit.txt 2> pit.err; then
    cat pit.err >&2
    echo "Error: could not read ${pits[0]}." >&2
    exit 1
fi
rm -f pit.err
grep -q "Partition Name:" pit.txt || { echo "Error: pit.txt has no partitions." >&2; exit 1; }
echo "Read ${pits[0]}, saved to pit.txt"

echo "== 4. Heimdall command =="
cmd=(heimdall flash)
userdata=()
skipped=()
while IFS='|' read -r part fname; do
    # The factory USERDATA image replaces the phone's data, which is the
    # factory reset: the CSC cache.img then makes recovery resize it to the
    # whole partition on first boot. Flash it last, so that if the transfer
    # fails everything else is already written. With --keep-data it must not
    # be flashed at all.
    if [ "$part" = "USERDATA" ]; then
        [ "$CSC_PREFIX" = "CSC_" ] && [ -f "extracted/$fname" ] && userdata=( "--$part" "extracted/$fname" )
        continue
    fi
    # The HIDDEN partition (mounted at /preload) only holds optional preloaded
    # content, e.g. a sample song on an S8. Some phones have a smaller HIDDEN
    # partition than the firmware's image: writing it overflows the partition
    # and the phone then shows "use the Emergency recovery function".
    [ "$part" = "HIDDEN" ] && continue
    [ -n "$fname" ] && [ -f "extracted/$fname" ] && cmd+=( "--$part" "extracted/$fname" )
done < <(awk -F': ' '
    /Partition Name:/ {part=$2}
    /Flash Filename:/ {print part "|" $2}' pit.txt)
cmd+=( "${userdata[@]}" )

for f in extracted/*; do
    name="${f##*/}"
    case "$name" in *.pit|meta-data) continue ;; esac
    printf '%s\n' "${cmd[@]}" | grep -qx "extracted/$name" || skipped+=("$name")
done

if [ "$CSC_PREFIX" = "CSC_" ] && [ "${#userdata[@]}" -eq 0 ]; then
    echo "Error: no userdata image in the AP archive, the phone would not be reset." >&2
    exit 1
fi

printf '%q ' "${cmd[@]}"; echo
if [ "${#skipped[@]}" -gt 0 ]; then
    echo
    echo "Files not flashed (no matching partition in the PIT, HIDDEN, or USERDATA with --keep-data):"
    printf '  %s\n' "${skipped[@]}"
fi