#!/bin/sh
# som-recover-nrf91 -- bring an Emcraft SOM-NRF9151 whose debug access is locked
# (APPROTECT) back to an installed, booting release, from a brintOS machine,
# over SWD through a J-Link: the station's `nrfutil device recover` plus the
# install that follows it, in one openocd session.
#
#   * The CTRL-AP's ERASEALL is the only way into a locked nRF91x1. It erases
#     the flash and the whole UICR, and the access port then stays open only
#     until the next pin, power or watchdog reset: a release whose firmware
#     opens it at boot (NCS's NRF_APPROTECT_USE_UICR, which 2.0-122's b0,
#     MCUboot and TF-M are built with) does so only when UICR.APPROTECT and UICR.SECUREAPPROTECT both read
#     0x50FA50FA, and an erased UICR reads 0xFFFFFFFF. So both words are written
#     right after the erase, before anything else, and the release goes in whole,
#     its UICR records (b0's key hashes) included.
#   * openocd's nrf91_recover resets the chip after the erase through the
#     CTRL-AP. That is a soft reset, which Nordic's own procedure calls for; it
#     does not lock the port again. Nothing else here resets the chip until the
#     image is written and verified, and that reset is a soft one too.
#   * The KMU's key-slot part of the UICR (from 0x400) reads back erased until
#     the chip is reset, so a UICR page written back is verified up to 0x400 in
#     the erase session, and whole once the chip has reset.
#   * UICR.ERASEPROTECT (0x00FF8030) is never written: once set, ERASEALL is
#     refused and a locked chip cannot be recovered from a debugger at all.
#
# Exit status: 0 PASS, 1 FAIL, 2 usage, 4 the board is in use by another
# command (som-lock: the lock som-program-nrf91, som-test-nrf91 and som-console
# share), so nothing was done. A recover has no stop request: once started, it
# runs to its end.

set -u

PROG=som-recover-nrf91
SCRIPTS="${SOM_OPENOCD_SCRIPTS:-/usr/share/openocd/scripts}"
BOARD_CFG="${SOM_BOARD_CFG:-$SCRIPTS/board/emcraft-som-nrf9151.cfg}"
FW_DIR=/usr/share/brintos-som/nrf9151
FLASH_SIZE=0x100000
CTRL_AP_IDR=0x12880000

image=""
app=""
uicr_file=""
speed=""
adapter=""
yes=0
force=0
dry_run=0
ocd_ran=0

usage() {
    printf '%s\n' \
"usage: $PROG (--app NAME | --image FILE.hex) [--uicr FILE] [--speed KHZ]" \
"                         [--adapter NAME] [--force] [--yes] [--dry-run]" \
"" \
"  --app NAME      the release to install after the erase: $(ls "$FW_DIR" 2>/dev/null | tr '\n' ' ')" \
"  --image FILE    an Intel HEX to install instead; its UICR records are written" \
"  --uicr FILE     a 4 KiB UICR page saved from this board (e.g. its golden backup)," \
"                  written back instead of the image's UICR records" \
"  --yes           erase: everything on the chip goes, the UICR included" \
"  --force         recover a chip whose debug access is open (it does not need" \
"                  it: som-program-nrf91 installs and keeps the UICR); its flash" \
"                  and UICR are backed up first" \
"  --dry-run       resolve and report, touch no hardware" \
"" \
"Backup (with --force): to \$SOM_BACKUP_DIR, default /tmp." \
"A run that fails keeps openocd's output in \${TMPDIR:-/tmp}/$PROG-<time>-<pid>." \
"Never written: UICR.ERASEPROTECT."
}

while [ $# -gt 0 ]; do
    case "$1" in
        --app)      app=${2:-}; shift 2 || exit 2 ;;
        --image)    image=${2:-}; shift 2 || exit 2 ;;
        --uicr)     uicr_file=${2:-}; shift 2 || exit 2 ;;
        --speed)    speed=${2:-}; shift 2 || exit 2 ;;
        --adapter)  adapter=${2:-}; shift 2 || exit 2 ;;
        --yes)      yes=1; shift ;;
        --force)    force=1; shift ;;
        --dry-run)  dry_run=1; shift ;;
        -h|--help)  usage; exit 0 ;;
        *)          echo "$PROG: unknown argument: $1" >&2; usage >&2; exit 2 ;;
    esac
done
# A chip this run has erased reads as open until its next reset, so a run again
# needs --force.
again="$PROG ${app:+--app $app}${image:+--image $image}${uicr_file:+ --uicr $uicr_file}${speed:+ --speed $speed}${adapter:+ --adapter $adapter} --force --yes"

T0=$(date +%s)
now()  { date +%s; }
took() { s=$(( $(now) - T0 )); echo "$((s / 60)) min $((s % 60)) s"; }
fail() { keep_logs; echo "$PROG: FAIL -- $* [$(took)]"; exit 1; }
say()  { echo "$PROG: $*"; }
# The exit removes the work directory; a failed run first moves out what
# openocd printed in each session, and the sessions.
keep_logs() {
    [ "$ocd_ran" -eq 1 ] || return 0
    k=${TMPDIR:-/tmp}/$PROG-$(date -u +%Y%m%dT%H%M%SZ)-$$
    if mkdir "$k" 2>/dev/null && mv "$W"/*.log "$W"/*.tcl "$k"/ 2>/dev/null; then
        say "logs:     openocd's output from this run is kept in $k"
    else
        say "logs:     could not keep openocd's output in $k"
    fi
}

[ -n "$app" ] && [ -n "$image" ] && fail "--app and --image both name an image; give one"
if [ -n "$app" ]; then
    case "$app" in */*|.*) fail "--app takes a release name: $(ls "$FW_DIR" 2>/dev/null | tr '\n' ' ')" ;; esac
    image=$FW_DIR/$app/$app.hex
    [ -r "$image" ] || fail "no release '$app' on this machine"
fi
[ -n "$image" ] || { usage >&2; exit 2; }
[ -r "$image" ] || fail "cannot read $image"
case "$image" in
    *.hex|*.HEX|*.ihex) ;;
    *) fail "$image is not an Intel HEX: a recovered chip needs the image's UICR records, which a .bin cannot carry" ;;
esac
sum=$(sha256sum "$image" | cut -d' ' -f1)
say "image:    $image ($sum)"

W=${TMPDIR:-/tmp}/$PROG.$$
mkdir -p "$W" || fail "cannot make $W"
SOM_LOCK=${SOM_LOCK:-/tmp/som-board.lock}
locked=0
lock_board() {   # without som-lock on PATH the run is unlocked
    command -v som-lock >/dev/null 2>&1 || return 0
    if l_why=$(som-lock take "$PROG" $$); then locked=1; return 0; fi
    echo "$PROG: FAIL -- ${l_why:-cannot take the board lock $SOM_LOCK}; nothing was done [$(took)]"
    exit 4
}
unlock_board() {
    [ "$locked" -eq 0 ] || som-lock release $$
    locked=0
}
trap 'rm -rf "$W"; unlock_board' EXIT

# The image's records, split: the flash part, and its UICR records checked so
# that none of them touches ERASEPROTECT. With --uicr the UICR comes from the
# file instead, and the image's own records are left out.
awk -v flash="$W/flash.hex" -v keepuicr="$([ -n "$uicr_file" ] && echo 0 || echo 1)" '
    function h(s) { return index("0123456789ABCDEF", toupper(substr(s, 1, 1))) * 16 - 16 + index("0123456789ABCDEF", toupper(substr(s, 2, 1))) - 1 }
    { sub(/\r$/, "") }
    $0 == "" { next }
    substr($0, 1, 1) != ":" { bad = "line " NR " is not a HEX record"; exit }
    {
        n = h(substr($0, 2, 2)); t = substr($0, 8, 2)
        if (t == "04") base = (h(substr($0, 10, 2)) * 256 + h(substr($0, 12, 2))) * 65536
        if (t == "00") {
            a = base + h(substr($0, 4, 2)) * 256 + h(substr($0, 6, 2))
            if (a >= 16711680) {
                if (a < 16744500 && a + n > 16744496) { bad = sprintf("a UICR record at 0x%X writes ERASEPROTECT", a); exit }
                nu += n
                if (!keepuicr) next
            } else if (a + n > 1048576) { bad = sprintf("data at 0x%X is outside the flash and the UICR", a); exit }
            else data += n
        }
        print > flash
    }
    END { if (bad != "") { print "BAD " bad; exit 1 }; printf "DATA %d\nUICR %d\n", data, nu }' "$image" >"$W/split.txt" ||
    fail "$image: $(sed -n 's/^BAD //p' "$W/split.txt")"
data=$(sed -n 's/^DATA //p' "$W/split.txt"); nu=$(sed -n 's/^UICR //p' "$W/split.txt")
[ "${data:-0}" -gt 0 ] || fail "$image has no data inside the application flash"

# The image's busybox od takes no -A/-t/-N, so each word is compared whole.
printf '\372\120\372\120' > "$W/open"; printf '\377\377\377\377' > "$W/unset"
word() { dd if="$1" of="$W/word" bs=4 skip=$2 count=1 2>/dev/null && cmp -s "$W/word" "$W/$3"; }
if [ -n "$uicr_file" ]; then
    [ "$(stat -c %s "$uicr_file" 2>/dev/null)" = 4096 ] || fail "$uicr_file is not a 4 KiB UICR page"
    word "$uicr_file" 0 open && word "$uicr_file" 11 open ||
        fail "$uicr_file does not open the access port (APPROTECT and SECUREAPPROTECT must both be 0x50FA50FA): the chip would lock at its next reset"
    word "$uicr_file" 12 unset || fail "$uicr_file sets ERASEPROTECT; that is never written"
    say "uicr:     to write back $uicr_file ($(sha256sum "$uicr_file" | cut -d' ' -f1)); the image's $nu UICR bytes left out"
else
    say "uicr:     to write APPROTECT and SECUREAPPROTECT = 0x50FA50FA, then the image's $nu UICR bytes"
fi
say "content:  $data bytes of flash"
say "erase:    ERASEALL through the CTRL-AP -- the flash and the UICR"
[ -r "$BOARD_CFG" ] || fail "no board config at $BOARD_CFG"
if [ "$dry_run" -eq 1 ]; then
    say "dry-run, nothing touched"
    say "PASS -- resolved $image [$(took)]"
    exit 0
fi
command -v openocd >/dev/null 2>&1 || fail "openocd is not installed"
lock_board

pre=""
[ -n "$adapter" ] && pre="$pre set BRINTOS_JTAG_IFACE interface/$adapter.cfg;"
[ -n "$speed" ] && pre="$pre set BRINTOS_JTAG_SPEED $speed;"
ocd() {   # $1 = phase Tcl, $2 = log
    ocd_ran=1
    if [ -n "$pre" ]; then openocd -c "$pre" -f "$BOARD_CFG" -f "$1" >"$2" 2>&1
    else openocd -f "$BOARD_CFG" -f "$1" >"$2" 2>&1; fi
}
show() { grep -E 'SOM |Error|VTarget' "$1" | sed 's/^/  | /'; }
mark() { sed -n "s/^SOM $1=//p" "$2" | sed -n 1p; }

# --- 1. the CTRL-AP: which chip, is it locked, may it be erased -----------------------
printf '%s\n' 'init' 'set dap nrf91.dap' \
    'echo [format "SOM idr=0x%08x" [$dap apreg 4 0xfc]]' \
    'echo [format "SOM approtect-status=%d" [$dap apreg 4 0x0c]]' \
    'echo [format "SOM eraseprotect-status=%d" [$dap apreg 4 0x18]]' \
    'shutdown' >"$W/probe.tcl"
ocd "$W/probe.tcl" "$W/probe.log"
show "$W/probe.log"
idr=$(mark idr "$W/probe.log"); aps=$(mark approtect-status "$W/probe.log"); eps=$(mark eraseprotect-status "$W/probe.log")
[ "$idr" = "$CTRL_AP_IDR" ] || fail "no nRF91x1 CTRL-AP (IDR '${idr:-none}', want $CTRL_AP_IDR): is the J-Link on the SOM's SWD header?; nothing was done"
[ $(( ${eps:-0} & 1 )) -eq 1 ] ||
    fail "ERASEPROTECT is set: the chip refuses ERASEALL and cannot be recovered from a debugger; nothing was done"
if [ "${aps:-0}" -ge 3 ]; then
    [ "$force" -eq 1 ] || fail "debug access is open (APPROTECT.STATUS $aps): this chip needs no recover -- som-program-nrf91 installs and keeps its UICR. --force erases it anyway; nothing was done"
    say "chip:     debug access open (APPROTECT.STATUS $aps); --force: backing up, then erasing"
else
    say "chip:     debug access locked (APPROTECT.STATUS $aps)"
fi
[ "$yes" -eq 1 ] || fail "this erases the whole chip, the UICR included; give --yes; nothing was done"

# --- 2. a chip that can still be read is read first (--force) --------------------------
if [ "${aps:-0}" -ge 3 ]; then
    bdir=${SOM_BACKUP_DIR:-/tmp}
    stem=$bdir/nrf91-recover-$(date -u +%Y%m%dT%H%M%SZ)
    printf '%s\n' 'init' 'set ok 1' \
        "if {[catch {dump_image $stem-flash.bin 0x0 $(( FLASH_SIZE ))} e]} { set ok 0; echo \"SOM backup=fail \$e\" }" \
        "if {\$ok && [catch {dump_image $stem-uicr.bin 0x00FF8000 4096} e]} { set ok 0; echo \"SOM backup=fail \$e\" }" \
        'if {$ok} { echo "SOM backup=pass" }' 'shutdown' >"$W/backup.tcl"
    ocd "$W/backup.tcl" "$W/backup.log"
    [ "$(mark backup "$W/backup.log")" = pass ] && [ "$(stat -c %s "$stem-flash.bin" 2>/dev/null)" = $(( FLASH_SIZE )) ] &&
        [ "$(stat -c %s "$stem-uicr.bin" 2>/dev/null)" = 4096 ] ||
        { show "$W/backup.log"; fail "could not read the chip out first; nothing was done"; }
    say "backup:   $stem-flash.bin ($(sha256sum "$stem-flash.bin" | cut -d' ' -f1)), $stem-uicr.bin ($(sha256sum "$stem-uicr.bin" | cut -d' ' -f1))"
fi

# --- 3. erase, open the port for good, write, verify -- one session ------------------
if [ -n "$uicr_file" ]; then
    dd if="$uicr_file" of="$W/uicr-head.bin" bs=1024 count=1 2>/dev/null || fail "cannot read $uicr_file"
    uicr_step="flash write_image $uicr_file 0x00FF8000 bin"
    uicr_verify="verify_image $W/uicr-head.bin 0x00FF8000 bin"
else
    uicr_step='flash fillw 0x00FF8000 0x50FA50FA 1; flash fillw 0x00FF802C 0x50FA50FA 1'
    uicr_verify='if {[r32 0x00FF8000] != 0x50FA50FA || [r32 0x00FF802C] != 0x50FA50FA} { error "APPROTECT words not written" }'
fi
printf '%s\n' \
    'proc r32 {a} { set s [mdw $a]; if {![regexp {:[ \t]*([0-9a-fA-F]+)} $s -> v]} { error "mdw $a: $s" }; return [expr 0x$v] }' \
    'init' 'set dap nrf91.dap' \
    'set rc [catch {nrf91_recover} e]; echo "SOM recover=$rc $e"' \
    'if {$rc == 0 && [$dap apreg 4 0x0c] < 3} { set rc 1; echo "SOM recover=1 the access port did not open" }' \
    'if {$rc == 0} { set rc [catch {halt} e]; echo "SOM halt=$rc $e" }' \
    "if {\$rc == 0} { set rc [catch {$uicr_step} e]; echo \"SOM uicr=\$rc \$e\" }" \
    "if {\$rc == 0} { set rc [catch {flash write_image $W/flash.hex} e]; echo \"SOM write=\$rc \$e\" }" \
    "if {\$rc == 0} { set rc [catch {verify_image $W/flash.hex} e]; echo \"SOM verify=\$rc \$e\" }" \
    "if {\$rc == 0} { set rc [catch {$uicr_verify} e]; echo \"SOM uicr-verify=\$rc \$e\" }" \
    'if {$rc == 0} { reset run; echo "SOM result=pass" } else { echo "SOM result=fail" }' \
    'shutdown' >"$W/recover.tcl"
say "recover:  erase, UICR, write, verify ..."
t1=$(now)
ocd "$W/recover.tcl" "$W/recover.log"
show "$W/recover.log"
# A retry session that som-program-nrf91 kept after a failed write describes the
# chip before this erase. Left in place it stops every later som-program-nrf91
# run, pointing at a retry that would write the failed run's image over this
# one, so once the erase has run each is set aside as som-nrf91-recovered-*.
case "$(mark recover "$W/recover.log")" in
    0*) for d in "${TMPDIR:-/tmp}"/som-nrf91-retry-*; do
            [ -d "$d" ] || continue
            mv "$d" "${TMPDIR:-/tmp}/som-nrf91-recovered-${d##*/som-nrf91-retry-}" &&
                say "retry:    set aside $d: a failed write's retry session from before this erase"
        done ;;
esac
if [ "$(mark result "$W/recover.log")" != pass ]; then
    case "$(mark recover "$W/recover.log")" in
        0*) fail "the chip was erased but the write did not complete; run $again now, before any reset: the chip stays unlocked only until its next pin, power or watchdog reset, and the erase left its debug port open, so the run needs --force" ;;
        *)  fail "ERASEALL did not complete; nothing was written" ;;
    esac
fi
say "written:  verified and reset ($(( $(now) - t1 )) s)"

# --- 4. what the chip says after its reset: the port, and the UICR it now reads ------
# Just after the reset the CPU may not answer the debugger yet, so both reads retry
# for up to 20 s, and the CPU is examined again before each UICR read. Reading it
# during the boot left the application stalled (1.2-114), so the chip is reset once more
# after the read, with no vector catch set.
printf '%s\n' 'init' 'set dap nrf91.dap' 'set s 0' \
    'for {set i 0} {$i < 20} {incr i} { set s [$dap apreg 4 0x0c]; if {$s >= 3} break; sleep 1000 }' \
    'echo [format "SOM approtect-status=%d" $s]' \
    'set rc 1; set e ""; set i 0' \
    "while {\$s >= 3 && \$rc && \$i < 20} { if {\$i} { sleep 1000 }; incr i; set rc [catch {nrf91.cpu arp_examine; dump_image $W/uicr-after.bin 0x00FF8000 4096} e] }" \
    'if {$s >= 3} { if {$rc} { echo "SOM uicr-read=fail $e" } else { echo "SOM uicr-read=pass"; echo "SOM uicr-tries=$i" } }' \
    'if {$s >= 3} { catch {cortex_m vector_catch none}; set r [catch {reset run} e]; echo "SOM rerun=$r $e" }' \
    'shutdown' >"$W/after.tcl"
ocd "$W/after.tcl" "$W/after.log"
aps=$(mark approtect-status "$W/after.log"); tries=$(mark uicr-tries "$W/after.log")
say "after:    APPROTECT.STATUS ${aps:-none} (3 = open)${tries:+; the UICR read on try $tries}"
[ -n "$aps" ] || { show "$W/after.log"; fail "could not reach the chip after its reset"; }
[ "$aps" -ge 3 ] || fail "the chip locked again at its reset: nothing on it can be read until another recover"
case "$(mark rerun "$W/after.log")" in
    0*) ;;
    *)  show "$W/after.log"; fail "could not reset the chip after reading its UICR; reset or power-cycle the board" ;;
esac
[ "$(mark uicr-read "$W/after.log")" = pass ] && [ "$(stat -c %s "$W/uicr-after.bin" 2>/dev/null)" = 4096 ] ||
    { show "$W/after.log"; fail "could not read the UICR back after the reset"; }
if [ -n "$uicr_file" ]; then
    d=$(cmp "$uicr_file" "$W/uicr-after.bin" 2>&1) ||
        fail "the UICR read back after the reset is not $uicr_file ($d); the chip runs with its port open: run $again"
    say "uicr:     read back after the reset: all 4096 bytes equal $uicr_file"
else
    word "$W/uicr-after.bin" 0 open && word "$W/uicr-after.bin" 11 open ||
        fail "the UICR read back after the reset does not hold APPROTECT and SECUREAPPROTECT = 0x50FA50FA"
    say "uicr:     read back after the reset: APPROTECT and SECUREAPPROTECT 0x50FA50FA"
fi
say "PASS -- recovered: erased, UICR APPROTECT 0x50FA50FA, wrote $image ($sum) [$(took)]"