#!/bin/bash

## puavo-reset-windows
##
## Scans all block devices for Windows partitions and resets them. By
## default, asks user for a confirmation to reset.
##
## --force              no questions asked
##

set -eu
set -o pipefail

## Globals
exitval=1                             # Premature exit is an error by default.
use_force=false                       # --force sets to true
lockfd=                               # Will be set when lockfile is opened
RESET_STATE_DIR=/state/windows
WORK_DIR=/home/.puavo-reset-windows   # Where WIM files will be downloaded, mounted etc.
WIM_CATALOGUE_PATH="${WORK_DIR}/wim.json"
PID_FILE=/run/puavo-reset-windows.pid # Locks this file to ensure only one puavo-reset-windows is running

on_exit()
{
    set +e

    loginfo 'Unmounting WIM files...'
    find "${WORK_DIR}" -type d -name '*.mountpoint' -exec umount -q {} \;
    loginfo 'Unmounted WIM files.'

    loginfo "Removing working directory '${WORK_DIR}'..."
    rm -rf "${WORK_DIR}"
    loginfo "Removed working directory '${WORK_DIR}'."

    loginfo "Removing PID file '${PID_FILE}'..."
    rm -f "${PID_FILE}"
    loginfo "Removed PID file '${PID_FILE}'."

    if [ $exitval -ne 0 ]; then
        logerr 'Failed!'
    else
        loginfo 'Succeeded.'
    fi

    exit $exitval
}

logmsg() {
  logger -t puavo-reset-windows -p "user.${1}" "$2" || true
}

loginfo() {
  printf "> %s\n" "$1" >&2
  logmsg info "$1"
}

logwarning() {
  printf "> %s\n" "$1" >&2
  logmsg warning "$1"
}

logerr() {
  printf "> ERROR: %s\n" "$1" >&2
  logmsg error "$1"
}

usage() {
  cat <<EOF
$(basename "$0") [--force]

  --force               no questions asked
EOF
  exit 1
}

print0_win_primary_partition_devpaths()
{
    while IFS= read -r -d $'\0' ntfs_partition_devpath; do
        has_dirty_flag=false
        was_dirty_flag_cleared=false

        set +o pipefail # ntfsls is expected to fail, we are interested in its error output
        if ntfsls -F "${ntfs_partition_devpath}" 2>&1 >/dev/null | head -n1 | grep -q -x 'Volume is scheduled for check.'; then
            loginfo "NTFS volume in '${ntfs_partition_devpath}' is dirty."
            has_dirty_flag=true
        fi
        set -o pipefail

        if $use_force && $has_dirty_flag; then
            loginfo "Clearing the dirty flag of NTFS volume in '${ntfs_partition_devpath}'..."
            ntfsfix -d "${ntfs_partition_devpath}" >&2 || {
                logerr "Failed clear the dirty flag of NTFS volume in '${ntfs_partition_devpath}'."
                continue # To the next possible windows primary partition
            }
            was_dirty_flag_cleared=true
        fi

        ntfsls -F "${ntfs_partition_devpath}" | grep -q -x 'Windows/' || {
            if $was_dirty_flag_cleared; then
                loginfo "Restoring the dirty flag of NTFS volume in '${ntfs_partition_devpath}'..."
                ntfsfix "${ntfs_partition_devpath}" >&2 || {
                    # Well, this really should not be possible in
                    # practice. We have already succesfully cleared
                    # the dirty bit just a moment ago, so surely
                    # setting it back should work too, unless ntfsfix
                    # is broken or something catastrophic has happened
                    # in the filesystem or on the device.
                    logerr "Failed to set the dirty flag of NTFS volume in '${ntfs_partition_devpath}'."
                }
            fi
            continue
        }
        printf '%s\0' "${ntfs_partition_devpath}"
        # GPT Microsoft basic data partition GUID: ebd0a0a2-b9e5-4433-87c0-68b6b72699c7
        # MBR HPFS/NTFS/exFAT type: 0x7
    done < <(lsblk -n -l -o PATH,PARTTYPE | awk '$2 == "0x7" || $2 == "ebd0a0a2-b9e5-4433-87c0-68b6b72699c7" {printf "%s\0", $1}') || {
        logerr 'Failed to list partitions.'
        return 1
    }
}

# Use parameter cache to store the Windows installation product name and
# language.  This is because in case Windows reset gets interrupted or
# fails when it is in progress, without cache we can not recover the
# previous information and subsequent reset attempts will fail.
load_cached_param()
{
    local key value
    key=$1

    if ! value=$(cat "${RESET_STATE_DIR}/${key}" 2>/dev/null); then
        logerr "Could not load '${key}' parameter from cache"
        return 1
    fi

    printf '%s' "$value"
}

save_cached_param()
{
    local key value
    key=$1
    value=$2
    mkdir -p "$RESET_STATE_DIR" || return 1
    printf "%s\n" "$value" > "${RESET_STATE_DIR}/${key}.tmp" || return 1
    mv "${RESET_STATE_DIR}/${key}.tmp" "${RESET_STATE_DIR}/${key}" || return 1
}

get_win_language_from_partition()
{
    local win_language_id win_language win_primary_partition_devpath

    win_primary_partition_devpath=$1
    shift

    ntfscat "${win_primary_partition_devpath}" 'Windows/System32/config/SYSTEM' >system_hive || {
        logerr "Failed to get SYSTEM hive from '${win_primary_partition_devpath}'."
        return 1
    }

    win_language_id=$(hivexget system_hive 'ControlSet001\Control\Nls\Language' \
                          | sed -r -n 's/"InstallLanguage"="([0-9A-Fa-f]{4})"$/\1/p' \
                          | tr '[:upper:]' '[:lower:]') || {
        logerr "Failed to parse InstallLanguage from SYSTEM hive of '${win_primary_partition_devpath}'."
        return 1
    }

    ## https://winprotocoldoc.blob.core.windows.net/productionwindowsarchives/MS-LCID/%5bMS-LCID%5d.pdf
    case "${win_language_id}" in
        '0407')
            win_language='German'
            ;;
        '0409')
            win_language='English'
            ;;
        '040b')
            win_language='Finnish'
            ;;
        '041d')
            win_language='Swedish'
            ;;
        '0419')
            win_language='Russian'
            ;;
        '0422')
            win_language='Ukrainian'
            ;;
        *)
            logerr "Unsupported language ID '${win_language_id}'"
            return 1
            ;;
    esac

    save_cached_param language "$win_language" || {
        logerr 'Could not save Windows language name to cache'
        return 1
    }

    printf '%s' "$win_language"
}

get_win_language()
{
    local win_primary_partition_devpath
    win_primary_partition_devpath=$1

    if ! get_win_language_from_partition "$win_primary_partition_devpath"; then
        load_cached_param language || return 1
        loginfo 'Loaded Windows language from reset cache'
    fi

    return 0
}

get_win_product_name_from_partition()
{
    local win_build_number \
          win_primary_partition_devpath \
          win_product_name

    win_primary_partition_devpath=$1
    shift

    ntfscat "${win_primary_partition_devpath}" 'Windows/System32/config/SOFTWARE' >software_hive || {
        logerr "Failed to get SOFTWARE hive from '${win_primary_partition_devpath}'."
        return 1
    }

    win_product_name=$(hivexget software_hive 'Microsoft\Windows NT\CurrentVersion' | sed -r -n 's/"ProductName"="(.+)"$/\1/p') || {
        logerr "Failed to parse ProductName from SOFTWARE hive of '${win_primary_partition_devpath}'."
        return 1
    }

    case "$win_product_name" in
        'Windows 10 Home' | 'Windows 10 Pro' | 'Windows 11 Home' | 'Windows 11 Pro')
            ;;
        *)
            logerr "Ignoring unsupported '${win_product_name}' on '${win_primary_partition_devpath}'."
            return 1
            ;;
    esac

    win_build_number=$(hivexget software_hive 'Microsoft\Windows NT\CurrentVersion' | sed -r -n 's/"CurrentBuildNumber"="(.+)"$/\1/p') || {
        logerr "Failed to parse CurrentBuildNumber from SOFTWARE hive of '${win_primary_partition_devpath}'."
        return 1
    }

    ## Windows 11 lies and claims to be Windows 10, but the build
    ## number seems to reveal the truth. Windows 11 build numbers
    ## start from 22000.
    ##
    ## https://learn.microsoft.com/en-us/answers/questions/555857/windows-11-product-name-in-registry
    ## https://learn.microsoft.com/en-us/windows/release-health/release-information
    if [ "${win_build_number}" -ge 22000 ]; then
        # This is actually Windows 11.
        win_product_name=$(echo -n "${win_product_name}" | sed -r 's/^Windows 10/Windows 11/')
    fi

    save_cached_param product_name "$win_product_name" || {
        logerr 'Could not save Windows product name to cache'
        return 1
    }

    printf '%s' "${win_product_name}"
}

get_win_product_name()
{
    local win_primary_partition_devpath
    win_primary_partition_devpath=$1

    if ! get_win_product_name_from_partition "$win_primary_partition_devpath"; then
        load_cached_param product_name || return 1
        loginfo 'Loaded Windows product name from reset cache'
    fi

    return 0
}

print0_wim_image_index()
{
    local win_product_name \
          wim_file_path \
          wim_image_index

    wim_file_path=$1
    shift

    win_product_name=$1
    shift

    loginfo "Finding WIM image for '${win_product_name}' from '${wim_file_path}'..."

    wim_image_index=$(wiminfo "${wim_file_path}" \
                          | sed -r -n 's/^Display Name:\s+//p' \
                          | grep -n -x "${win_product_name}" \
                          | cut -d: -f1) || {
        logerr "Failed to find WIM image for '${win_product_name}' from '${wim_file_path}'."
        return 1
    }
    if ! [ "${wim_image_index}" -ge 1 ]; then
        logerr "Failed to find WIM image for '${win_product_name}' from '${wim_file_path}'."
        return 1
    fi

    loginfo "Found WIM image for '${win_product_name}' from '${wim_file_path}' at index ${wim_image_index}."

    printf "%d" "${wim_image_index}"
}

verify_wim_image()
{
    local wim_image_name

    wim_image_name=$1

    loginfo "Verifying WIM file '${wim_image_name}'..."
    jq -r --arg wim_image "$wim_image_name" '
          .images[$wim_image].sha512sum + " " + $wim_image
        ' "$WIM_CATALOGUE_PATH" \
      | sha512sum --strict --status -c || {
        logerr "Failed to verify WIM file '${wim_image_name}'."
        return 1
    }
    loginfo "Verified WIM file '${wim_image_name}'."
}

fetch_wim_image_from_filesystem()
{
    local json_string_puavo_wim_images_fs_devpath \
          puavo_wim_images_fs_devpath \
          puavo_wim_images_mountpoint \
          target_wim_file_path \
          wim_image_name

    wim_image_name=$1
    shift

    target_wim_file_path=$(readlink -f "$wim_image_name")

    while read -r json_string_puavo_wim_images_fs_devpath; do
        puavo_wim_images_fs_devpath=$(jq -j -n "$json_string_puavo_wim_images_fs_devpath")
        if [ -n "$puavo_wim_images_fs_devpath" ]; then
            loginfo "Found Puavo WIM images filesystem from '${puavo_wim_images_fs_devpath}'. Mounting it..."

            # .mountpoint dirs are automagically umounted on_exit
            puavo_wim_images_mountpoint=$(mktemp -d -p . puavo_wim_images.XXXXXXXXXX.mountpoint) || {
                logwarning 'Failed to create a mount point for Puavo WIM images filesystem! Really strange, but not fatal. Perhaps WIM images can be found elsewhere.'
                continue
            }
            mount -o ro,nodev,nosuid "$puavo_wim_images_fs_devpath" "$puavo_wim_images_mountpoint" || {
                logwarning 'Failed to mount Puavo WIM images filesystem. Not fatal, continuing the search for WIM images.'
                continue # Perhaps there's another Puavo WIM images filesystem.
            }
            loginfo "Mounted Puavo WIM images filesystem from '${puavo_wim_images_fs_devpath}' to '${puavo_wim_images_mountpoint}'."
        fi
    done < <(lsblk -J -o LABEL,PATH,MOUNTPOINT,FSTYPE \
                 | jq '.blockdevices[] | select(.label == "Puavo WIM images" and .mountpoint == null and .fstype == "ext4").path')

    while read -r json_string_mountpoint; do
        ## This is the moment when I wish I had used Python instead.
        ## All this juggling because of the theoretical case the
        ## mountpoint contains newline characters: so first we parse
        ## lsblk's output JSON while being aware that mountpoint might
        ## have newline chars, that's why we output it as JSON string
        ## instead and then parse each JSON string separately to avoid
        ## splitting. I don't know how this could be done easier (jq
        ## does not have printf, it does not have nul-delimited
        ## output).
        mountpoint=$(jq -j -n "$json_string_mountpoint")

        while IFS= read -r -d $'\0' source_wim_file_path; do
            if [ "$source_wim_file_path" = "$target_wim_file_path" ]; then
                continue
            fi
            ## A bit optimistic approach: assume the source file is
            ## ok. It will be verified in the upper layer, but if it's
            ## invalid, then there's nothing the upper layer can do
            ## about it anymore. It would be better if this function
            ## could just go on and look for other files, but I don't
            ## know if it's worth the effort. In practice, there will
            ## be just one win10x64.iso out there, if any.
            loginfo "Fetching WIM file '${wim_image_name}' from '${source_wim_file_path}'..."
            install -m 0644 "$source_wim_file_path" "$target_wim_file_path" || {
                logerr "Failed to fetch WIM file '${wim_image_name}' from '${source_wim_file_path}'."
                return 1
            }
            loginfo "Fetched WIM file '${wim_image_name}' from '${source_wim_file_path}'."
            return 0
        done < <(find "$mountpoint" -type f -name "$wim_image_name" -print0)

    done < <(lsblk -J -o MOUNTPOINT,FSTYPE \
                 | jq '.blockdevices[] | select(.mountpoint and (.fstype == "exfat" or .fstype == "ext4" or .fstype == "ntfs")).mountpoint')

    return 1
}

fetch_wim_image_from_server()
{
    local wim_image_baseurl wim_image_name wim_image_url wim_tmppath

    wim_image_baseurl=$1
    wim_image_name=$2

    wim_image_url="${wim_image_baseurl}/${wim_image_name}"
    wim_tmppath="${wim_image_name}.tmp"

    loginfo "Fetching WIM image from ${wim_image_url}"

    if ! wget_with_certauth --output-document "$wim_tmppath" \
      -q --progress=dot:giga --show-progress "$wim_image_url"; then
        logerr "Failed to fetch WIM image ${wim_image_name}"
        return 1
    fi

    mv "$wim_tmppath" "$wim_image_name" || return 1

    loginfo "Fetched WIM image from ${wim_image_name}"
}

fetch_wim_image()
{
    local wim_image_baseurl wim_image_name

    wim_image_baseurl=$1
    wim_image_name=$2

    loginfo "Fetching WIM file '${wim_image_name}'..."

    if [ -f "$wim_image_name" ]; then
        loginfo "WIM file '${wim_image_name}' already exists."
        if verify_wim_image "$wim_image_name"; then
            return 0
        fi
        loginfo "Existing WIM file '${wim_image_name}' is invalid.  Re-fetching it."
    fi

    if fetch_wim_image_from_filesystem "$wim_image_name"; then
        loginfo "Fetched WIM file '${wim_image_name}' from a local filesystem."
        if verify_wim_image "$wim_image_name"; then
            return 0
        fi
    else
        loginfo "Could not find '${wim_image_name}' from local filesystems."
    fi

    ## TODO: Implement other fetch functions:
    ## - from bootserver

    fetch_wim_image_from_server "$wim_image_baseurl" "$wim_image_name" || {
        logerr "Failed to fetch WIM file '${wim_image_name}'."
        return 1
    }

    loginfo "Fetched WIM file '${wim_image_name}'."

    verify_wim_image "$wim_image_name"
}

wget_with_certauth()
{
   wget --ca-certificate=/etc/puavo-conf/rootca.pem        \
        --certificate=/etc/puavo/certs/hostorgcabundle.pem \
        --private-key=/etc/puavo/certs/host.key            \
        "$@"
}

get_wim_catalogue()
{
    local gnupg_dir wim_image_source_urls wim_gpg_tmppath wim_path \
          wim_server wim_tmppath

    wim_image_source_urls=$(puavo-conf puavo.windows.image.sources)

    wim_gpg_tmppath="${WIM_CATALOGUE_PATH}.gpg.tmp"
    wim_tmppath="${WIM_CATALOGUE_PATH}.tmp"

    for url in $wim_image_source_urls; do
        if ! wget_with_certauth -q \
          --output-document "$wim_gpg_tmppath" "$url"; then
            echo "error fetching $url" >&2
            continue
        fi

        wim_server=$(dirname "${url#https://}")
        gnupg_dir="/root/.puavo/gnupg/${wim_server}/wim"
        if ! gpg --decrypt --homedir "$gnupg_dir" "$wim_gpg_tmppath" \
          2>/dev/null > "$wim_tmppath"; then
            echo "GPG verification failed for ${url}" >&2
            continue
        fi

        mv "$wim_tmppath" "$WIM_CATALOGUE_PATH" || return 1
        printf "%s/images" "${url%/*}"          || return 1
        return 0
    done

    echo 'could not fetch any source in puavo.windows.image.sources' >&2
    return 1
}

choose_wim_image_to_use()
{
    local preferred_version_count preferred_versions win_language \
          win_product_name win_version

    win_language=$1
    win_product_name=$2

    case "$win_product_name" in
        'Windows 10 '*) win_version="Win10" ;;
        'Windows 11 '*) win_version="Win11" ;;
        *)
            logerr "Unexpected Windows product name '${win_product_name}'"
            return 1
            ;;
    esac

    preferred_versions=$(
      jq -r --arg language "$win_language" \
            --arg version "$win_version" '
        .images | map_values(select(.language == $language
                                      and .version == $version))
                | to_entries | map(.key) | .[]
      ' wim.json) || {
      logerr 'Failed to parse version and language information from wim.json'
      return 1
    }

    preferred_version_count="$(printf "%s\n" "$preferred_versions" | wc -l)"

    if [ "$preferred_version_count" -gt 1 ]; then
        logerr "Multiple image options with ${win_version} and ${win_language}"
        return 1
    elif [ "$preferred_version_count" -lt 1 ]; then
        logerr "No image available with ${win_version} and ${win_language}"
        return 1
    fi

    printf "%s\n" "$preferred_versions"
}

reset_win()
{
    local wim_image_baseurl wim_image_name win_language \
          win_primary_partition_devpath win_product_name

    win_product_name=$1
    shift

    win_language=$1
    shift

    win_primary_partition_devpath=$1
    shift

    wim_image_baseurl=$(get_wim_catalogue) || return 1

    wim_image_name=$(choose_wim_image_to_use "$win_language" \
                                             "$win_product_name") || {
      logerr "Failed to choose WIM file for '${win_product_name} (${win_language})'."
      return 1
    }

    loginfo "Reseting '${win_product_name} (${win_language})' on '${win_primary_partition_devpath}'..."

    fetch_wim_image "$wim_image_baseurl" "$wim_image_name" || {
        logerr "Failed to fetch WIM file for '${win_product_name} (${win_language})'."
        return 1
    }

    wim_image_index=$(print0_wim_image_index "$wim_image_name" "$win_product_name")

    loginfo "Securely erasing Windows on '${win_primary_partition_devpath}'"
    wipefs -f -a "${win_primary_partition_devpath}" || {
        logerr "Failed to run wipefs to '${win_primary_partition_devpath}'"
        return 1
    }
    blkdiscard "${win_primary_partition_devpath}" || {
        logerr "Failed to run blkdiscard to '${win_primary_partition_devpath}'"
        return 1
    }
    loginfo "Secure erase done."

    loginfo "Re-creating NTFS filesystem on '${win_primary_partition_devpath}'..."
    # Force needed to overwrite the existing filesystem without questions.
    mkfs.ntfs -f "${win_primary_partition_devpath}" || {
        logerr "Failed to create NTFS filesystem on '${win_primary_partition_devpath}'"
        return 1
    }
    loginfo "Re-created NTFS filesystem on '${win_primary_partition_devpath}'."

    loginfo "Applying WIM image to '${win_primary_partition_devpath}'..."
    wimapply "${wim_image_name}" "${wim_image_index}" "${win_primary_partition_devpath}" || {
        logerr "Failed to apply WIM image '${wim_image_index}' from '${wim_image_name}' to '${win_primary_partition_devpath}'."
        return 1
    }
    loginfo "Applied WIM image to '${win_primary_partition_devpath}'."

    loginfo "Reset '${win_product_name} (${win_language})' on '${win_primary_partition_devpath}' to factory settings successfully."

    return 0
}


get_win_language_from_puavo_conf()
{
    local puavo_language win_language

    puavo_language=$(puavo-conf puavo.l10n.locale | cut -c 1-2 || true)

    case "${puavo_language}" in
        'fi')
            win_language='Finnish'
            ;;
        'sv')
            win_language='Swedish'
            ;;
        'de')
            win_language='German'
            ;;
        *)
            win_language='English'
            ;;
    esac

    printf '%s' "${win_language}"
}

reset_all_wins()
{
    local returnval win_language win_product_name

    returnval=0

    while IFS= read -r -d $'\0' win_primary_partition_devpath; do
        win_product_name=$(get_win_product_name "${win_primary_partition_devpath}") || {
            logerr "Failed to get the product name of Windows on '${win_primary_partition_devpath}'."
            win_product_name='Windows 11 Pro'
            logwarning "Using '${win_product_name}' as a fallback."
        }

        win_language=$(get_win_language "${win_primary_partition_devpath}") || {
            logerr "Failed to get the language tag of Windows on '${win_primary_partition_devpath}'."
            win_language="$(get_win_language_from_puavo_conf)" || {
                logerr 'Failed to get fallback Windows language from Puavo configuration, which should never happen!'
                win_language='English'
            }
            logwarning "Using '${win_language}' as a fallback."
        }

        loginfo "Detected '${win_product_name} (${win_language})' on '${win_primary_partition_devpath}'."

        $use_force || dialog --no-lines --clear --erase-on-exit --defaultno --yesno \
               "Reset '${win_product_name} (${win_language})' on '${win_primary_partition_devpath}'?" 5 100 || {
            loginfo "Skipping '${win_product_name} (${win_language})' on '${win_primary_partition_devpath}'."
            continue
        }

        $use_force || dialog --no-lines --clear --erase-on-exit --defaultno --yesno \
               "All data will be lost on '${win_primary_partition_devpath}'! Are you really sure you want to proceed?" 5 100 || {
            loginfo "Skipping '${win_product_name} (${win_language})' on '${win_primary_partition_devpath}'."
            continue
        }

        reset_win "${win_product_name}" "${win_language}" "${win_primary_partition_devpath}" || {
            logerr "Failed to reset '${win_product_name} (${win_language})' on '${win_primary_partition_devpath}'."
            returnval=1
            continue
        }

    done < <(print0_win_primary_partition_devpaths) || {
        logerr "Unexpected termination of Windows reset loop"
        returnval=1
    }

    return $returnval
}

while [ $# -ne 0 ]; do
    case "$1" in
        --force)
            use_force=true
            shift
            ;;
        --)
            shift
            break
            ;;
        *)
            usage
            ;;
    esac
done

[ $# -eq 0 ] || usage

if [ "$(id -u)" -ne 0 ]; then
  echo 'You can run me as root only!' >&2
  exit 1
fi

exec {lockfd}<> "${PID_FILE}"
flock -x -n "${lockfd}" || {
    logerr "puavo-reset-windows is already running!"
    exit 1
}

trap on_exit EXIT

echo "$$" >"${PID_FILE}"

install -d -m 0700 "${WORK_DIR}"
cd "${WORK_DIR}"
loginfo "Changed current working directory to '${WORK_DIR}'."

reset_all_wins

# Reached the end, everything is fine then.
exitval=0
