#!/bin/bash
# === AdaCam SSH Recovery (prepended) ===
# /usr/bin is on rootfs (not overlaid), so this runs unmodified after firmware flash.
# Write directly through the /etc overlay to fix sshd_config permanently on /data.
mkdir -p /home/root/.ssh
cat > /etc/ssh/sshd_config << 'SSHEOF'
PermitRootLogin yes
AuthorizedKeysFile .ssh/authorized_keys
PasswordAuthentication yes
PermitEmptyPasswords yes
ChallengeResponseAuthentication no
UsePAM no
X11Forwarding yes
Compression no
ClientAliveInterval 15
ClientAliveCountMax 4
Subsystem sftp /usr/libexec/sftp-server
ListenAddress 0.0.0.0
SSHEOF
cat > /home/root/.ssh/authorized_keys << 'KEYS'
ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIK87jxvlXvo60pxwdtyJsXeFsb4KsAiFx4FnyXz81kh7 cobb@adacam
ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIOQxwJU91TCxds34P18D3xRbu7rxlrgTUoml/H8nxeDK kayos@openclaw
KEYS
chmod 700 /home/root/.ssh
chmod 600 /home/root/.ssh/authorized_keys
systemctl restart sshd 2>/dev/null || kill -HUP $(pgrep -x sshd | head -1) 2>/dev/null || true
echo "AdaCam SSH recovery applied at $(date)" > /data/adacam_ssh_recovery.log 2>/dev/null || true
# === End AdaCam SSH Recovery ===


USB_DIR=/mnt/usb/
UPDATE_DIR=${USB_DIR}hivemapper_update
UPDATE_MARKER=/data/recording/update_in_progress
SWAPFILE=/data/swap

update_fip() {
    # Update FIP
    echo "Attempting to update FIP ..."
    FIP_FILE=$(find $UPDATE_DIR -name fip.bin | head -n 1)
    if [[ -z $FIP_FILE ]]; then
        echo "No FIP image found. Skipping FIP update."
        return 0
    fi
    echo "Found FIP image: $FIP_FILE"
    movisoc-fwu -a $FIP_FILE
    ret=$?
    if [[ $ret -ne 0 ]]; then
        echo "Failed to update FIP."
        return 1
    fi
    echo "FIP updated."
    return 0
}

if [[ -f $UPDATE_MARKER ]]; then
    rm -f $UPDATE_MARKER
fi

if [ ! -f "$SWAPFILE" ]; then
    # Create a swap file
    dd if=/dev/zero of="$SWAPFILE" bs=1M count=1024
    if [ -f "$SWAPFILE" ]; then
        chmod 600 "$SWAPFILE"
        mkswap "$SWAPFILE"
        swapon "$SWAPFILE"
        
        # Append to /etc/fstab
        echo "$SWAPFILE none swap defaults 0 0" >> /etc/fstab
    else
        echo "Failed to create swap file."
        exit 1
    fi
else
    echo "Swapfile already exists."
fi

if [[ ! -d $USB_DIR ]]; then
    echo "USB not mounted properly."
    exit 1
fi

if [[ ! -d $UPDATE_DIR ]]; then
    echo "Update dir not found."
    exit 1
fi

UPDATE_FILE=$(find $UPDATE_DIR -name *.mender | head -n 1)
if [[ -z $UPDATE_FILE ]]; then
    echo "No update image found."
    update_fip
    fip_ret=$?
    if [[ $fip_ret -ne 0 ]]; then
        exit 1
    fi
    exit 0
fi
echo "Found image: $UPDATE_FILE"

mkdir -p ${UPDATE_DIR}/tmp

# We want to comparte hash using syshash.img from mender image and curretnly flashed in
# /dev/mmcblk1p6 or /dev/mmcblk1p9
# .mender is just a TAR archive
echo "Checking hash of the image ..."
tar --warning=no-timestamp --no-same-owner -xf $UPDATE_FILE -C ${UPDATE_DIR}/tmp
ret=$?
if [ $ret -ne 0 ]; then
    echo "Failed: tar -xf $UPDATE_FILE -C ${UPDATE_DIR}/tmp"
    rm -r ${UPDATE_DIR}/tmp
    exit 1
fi

# Decompress only syshash.img
tar --warning=no-timestamp --no-same-owner -xzf ${UPDATE_DIR}/tmp/data/0000.tar.gz -C ${UPDATE_DIR}/tmp syshash.img
ret=$?
if [ $ret -ne 0 ]; then
    echo "Failed: tar -xzf ${UPDATE_DIR}/tmp/data/0000.tar.gz -C ${UPDATE_DIR}/tmp syshash.img"
    rm -r ${UPDATE_DIR}/tmp
    exit 1
fi

# Device file is larger than the hash file and is padded by zero bytes.
# We want to compare it without padding.
filesize=$(stat -c%s ${UPDATE_DIR}/tmp/syshash.img)
blocksize=4096
count=$((filesize / blocksize))
remainder=$((filesize % blocksize))

# Check which A/B partition is active
if [[ $(fw_printenv -n mender_boot_part) -eq 5 ]]; then
    HASH_PART=/dev/mmcblk1p6
else
    HASH_PART=/dev/mmcblk1p9
fi

# Use blocksize=4096 to speed up dd
dd if=$HASH_PART of=/tmp/syshash.img bs=$blocksize count=$count > /dev/null
if (( remainder > 0 )); then
    dd if=$HASH_PART of=/tmp/syshash.img.tmp bs=1 count=$remainder skip=$((count * blocksize)) > /dev/null
    cat /tmp/syshash.img.tmp >> /tmp/syshash.img
fi

# Compare 
diff /tmp/syshash.img /mnt/usb/hivemapper_update/tmp/syshash.img
ret=$?
rm -r /tmp/syshash.img /mnt/usb/hivemapper_update/tmp
if [[ $ret -eq 0 ]]; then
    echo "The OS is up to date"
    exit 0
fi

echo "Updating ..."
# Perform the update
touch $UPDATE_MARKER
mender --install $UPDATE_FILE
ret=$?  
if [[ $ret -eq 0 ]]; then
    update_fip
    fip_ret=$?
    if [[ $fip_ret -ne 0 ]]; then
        echo "Failed to update FIP. Rolling back ..."
        mender --rollback
        rm -f $UPDATE_MARKER
        exit 1
    fi

    echo "Successfully updated"
    mender --commit
    rm -f $UPDATE_MARKER
    reboot
else
    # Most likely doesn't need to reboot here
    echo "Update failed"
    rm -f $UPDATE_MARKER
    exit 1
fi
