github.com/canonical/ubuntu-image@v0.0.0-20240430122802-2202fe98b290/tests/lib/external/snapd-testing-tools/remote/remote.refresh

#!/bin/bash -e

show_help() {
    echo "usage: remote.refresh snap [--channel CHANNEL] <SNAPNAME>"
    echo "usage: remote.refresh full [--channel CHANNEL]"
    echo "usage: remote.refresh disable-refreshes"
    echo ""
    echo "SNAPNAME: allowed options are: kernel, core, base, gadget, snapd or the snap name"
    echo ""
    echo "Available options:"
    echo "  -h --help   show this help message."
    echo ""
}

check_refresh_after_reboot() {
    local refresh_channel=$1
    local snap=$2

    echo "remote.refresh: checking snap \"$snap\" refresh after reboot"
    remote.wait-for reboot
    remote.wait-for snap-command
    remote.retry -n 10 --wait 2  "snap info $snap | grep -q -E \"tracking: +(latest/${refresh_channel}|${refresh_channel})\""
    echo "remote.refresh: snap \"$snap\" refreshed correctly"
}

check_refresh() {
    local refresh_channel=$1
    local snap=$2

    echo "remote.refresh: checking snap \"$snap\" refresh with no reboot"
    remote.wait-for ssh -n 60 --wait 2
    remote.wait-for snap-command
    remote.retry -n 10 --wait 2  "snap info $snap | grep -q -E \"tracking: +(latest/${refresh_channel}|${refresh_channel})\""
    echo "remote.refresh: snap \"$snap\" refreshed correctly"
}

check_waiting_for_reboot() {
    remote.exec "sudo journalctl -n 30" | grep -q "Waiting for system reboot"
}

check_ready_to_refresh() {
    remote.wait-for ssh
    if check_waiting_for_reboot; then
        echo "remote.refresh: waiting for system reboot"
        remote.wait-for reboot
    fi
    remote.wait-for refresh
}

do_refresh() {
    local snap_name=$1
    local refresh_channel=$2

    echo "remote.refresh: triggering refresh for snap \"$snap_name\" on channel \"$refresh_channel\""
    output=$(remote.exec "sudo snap refresh --channel ${refresh_channel} $snap_name 2>&1" || echo "snapd is about to reboot the system")
    if echo "$output" | grep -E "(no updates available|cannot refresh \"$snap_name\"|is not installed)"; then
        echo "remote.refresh: snap \"$snap_name\" has no updates available"
    elif echo "$output" | grep -E "snapd is about to reboot the system"; then
        remote.exec --timeout 3 "sudo reboot" || true
        check_refresh_after_reboot "$refresh_channel" "$snap_name"
    else
        check_refresh "$refresh_channel" "$snap_name"
    fi
}

process_refresh() {
    local snap_name=$1
    local refresh_channel=$2
    
    if [ -z "$refresh_channel" ]; then
        # Tracking is retrieved from snap info command because in old versions of
        # snapd the tracking is not included in the snap list output
        refresh_channel="$(remote.exec "snap info $snap_name" | grep -E '^tracking:' |  awk '{ print $2 }')"
    fi

    echo "remote.refresh: Refreshing $snap_name snap from $refresh_channel channel"
    do_refresh "$snap_name" "$refresh_channel"
}

refresh_fundamental() {
    snap_name=$1
    regex=$2
    refresh_channel=$3

    echo "remote.refresh: starting $snap_name refresh process"
    snap_name="$(remote.exec "snap list" | grep -E "$regex" | awk '{ print $1 }')"
    if [ -z "$snap_name" ]; then
        echo "remote.refresh: no $snap_name snap to update"
        return
    fi

    if [ "$(echo "$snap_name" | wc -l)" -gt 1 ]; then
        echo "remote.refresh: there is more than 1 $snap_name snap to update, skipping"
        return 0
    fi
    process_refresh "$snap_name" "$refresh_channel"
}

refresh_kernel() {
    local refresh_channel=$1
    local snap_name regex

    snap_name='kernel'
    regex='(kernel$|kernel,)'
    refresh_fundamental "$snap_name" "$regex" "$refresh_channel"
}

refresh_gadget() {
    local refresh_channel=$1
    local snap_name regex

    snap_name='gadget'
    regex='(gadget$|gadget,)'
    refresh_fundamental "$snap_name" "$regex" "$refresh_channel"
}

refresh_snapd() {
    local refresh_channel=$1
    local snap_name regex

    snap_name='snapd'
    regex='^snapd.*(snapd$|snapd,)'
    refresh_fundamental "$snap_name" "$regex" "$refresh_channel"
}

refresh_core() {
    local refresh_channel=$1
    local snap_name regex

    snap_name='core'
    regex='^core.*(core$|core,)'
    refresh_fundamental "$snap_name" "$regex" "$refresh_channel"
}

refresh_core_base() {
    local refresh_channel=$1
    local snap_name regex

    snap_name='core base'
    regex='^core.* (base$|base,)'
    refresh_fundamental "$snap_name" "$regex" "$refresh_channel"
}

refresh_snap() {
    local snapname=$1
    local refresh_channel=$2

    if [ -z "$snapname" ]; then
        echo "remote.refresh: snap name to refresh is not specified"
        return 1
    fi

    if ! remote.exec "snap list $snapname"; then
        echo "remote.refresh: no $snapname snap to update"
        return 0
    fi

    snap_line="$(remote.exec "snap list $snapname" | tail -1)"
    process_refresh "$snap_line" "$refresh_channel"
}

refresh_all() {
    # Run update and make "|| true" to continue when the connection is closed by remote host or not any snap to update
    remote.exec "sudo snap refresh" || true
    remote.wait-for ssh
}


get_boot_id() {
    remote.exec "cat /proc/sys/kernel/random/boot_id"
}

prevent_autorefresh() {
    remote.wait-for ssh
    state_path="$(remote.exec 'find /writable -name state.json 2>/dev/null' || true)"
    if [ -z "$state_path" ]; then
        echo "remote.refresh: state file not found in writable"
        state_path="$(remote.exec 'find / -name state.json 2>/dev/null' || true)"
    fi
    if [ -z "$state_path" ]; then
        echo "remote.refresh: state.json file not found"
        exit 1
    fi

    remote.exec "sudo cp $state_path /tmp/state.json"
    remote.exec "sudo chmod 644 /tmp/state.json"
    remote.pull /tmp/state.json state.json
    echo "remote.refresh: state file retrieved"

    jq ".data[\"last-refresh\"] = \"$(date +%Y-%m-%dT%H:%M:%S%:z)\"" state.json > state.json.new
    echo "remote.refresh: state file just updated"

    remote.push state.json.new /tmp/state.json.new
    remote.exec "sudo cp /tmp/state.json.new $state_path"
    remote.exec "sudo chmod 600 $state_path"
    echo "remote.refresh: updated state file restored"
    rm -f state.json state.json.new
}

disable_refreshes() {
    if ! command -v jq &>/dev/null; then
        snap install --devmode jq
    fi

    echo "remote.refresh: modifying state to make it look like the last refresh just happened"
    remote.exec "sudo systemctl stop snapd.socket snapd.service"
    prevent_autorefresh
    remote.exec "sudo systemctl start snapd.socket snapd.service"

    echo "remote.refresh: minimizing risk of hitting refresh schedule"
    remote.exec "sudo snap set core refresh.schedule=00:00-23:59"
    remote.exec "sudo snap refresh --time --abs-time" | grep -Eq "last: 2[0-9]{3}"

}

full_refresh() {
    echo "remote.refresh: starting full refresh"

    check_ready_to_refresh
    disable_refreshes

    check_ready_to_refresh
    refresh_core

    check_ready_to_refresh
    refresh_core_base

    check_ready_to_refresh
    refresh_snapd

    check_ready_to_refresh
    refresh_kernel

    check_ready_to_refresh
    refresh_all

    check_ready_to_refresh
}

snap_refresh() {
    local channel snapname

    while [ $# -gt 0 ]; do
        case "$1" in
            --channel)
                channel=$2
                shift 2
                ;;
            *)
                snapname=$1
                shift
                ;;
        esac
    done

    if [ -z "$snapname" ]; then
        echo "remote.refresh: snap name to refresh is not specified"
        return 1
    fi

    case "$snapname" in
        core)
            refresh_core "$channel"
            ;;
        base)
            refresh_core_base "$channel"
            ;;
        snapd)
            refresh_snapd "$channel"
            ;;
        kernel)
            refresh_kernel "$channel"
            ;;
        gadget)
            refresh_gadget "$channel"
            ;;
        *)
            refresh_snap "$snapname" "$channel"
            ;;
    esac
}

main() {
    if [ $# -eq 0 ]; then
        show_help
        exit
    fi

    local action
    case "$1" in
        -h|--help)
            show_help
            exit
            ;;
        full)
            action=full_refresh
            shift
            ;;
        snap)
            action=snap_refresh
            shift
            ;;
        disable-refreshes)
            action=disable_refreshes
            shift
            ;;
        *)
            echo "remote.refresh: unsupported parameter $1" >&2
            exit 1
            ;;
    esac

    if [ -z "$(declare -f "$action")" ]; then
        echo "remote.refresh: no such command: $action"
        show_help
        exit 1
    fi

    "$action" "$@"
}

main "$@"