#!/sbin/sh

OUTFD="/proc/self/fd/$2"
ZIPFILE="$3"
TARGET_DEVICES=aurora
READONLY_BLOCKS=
BLOCK_IDS=
STAGED_IMAGE="/tmp/xffc-image-$$"
umask 077

ui_print() {
  printf 'ui_print %s\nui_print\n' "$1" >> "$OUTFD"
}

abort() {
  ui_print "$1"
  exit 1
}

require_tools() {
  for tool in blockdev dd getprop head readlink rm sed sha256sum stat tr unzip; do
    command -v "$tool" >/dev/null 2>&1 || abort "[ERROR] Missing required command: $tool"
  done
  (set -o pipefail) 2>/dev/null || abort "[ERROR] Recovery shell does not support pipefail"
  set -o pipefail

  blockdev_help=$(blockdev --help 2>&1 || :)
  for option in --getsize64 --getro --setrw --setro; do
    case "$blockdev_help" in *"$option"*) ;; *)
      abort "[ERROR] blockdev does not support $option" ;;
    esac
  done

  probe="/tmp/xffc-dd-probe-$$"
  printf x | dd of="$probe" bs=1 conv=fsync >/dev/null 2>&1 || {
    rm -f "$probe"
    abort "[ERROR] dd does not support conv=fsync"
  }
  [ "$(dd if="$probe" iflag=count_bytes count=1 2>/dev/null)" = x ] || {
    rm -f "$probe"
    abort "[ERROR] dd does not support byte-count reads"
  }
  rm -f "$probe"
}

is_target_device() {
  value="$1"
  for target in $TARGET_DEVICES; do
    [ "$value" = "$target" ] && return 0
  done
  return 1
}

is_target_device_list() {
  values=$(printf '%s' "$1" | tr '|,:' '\n\n\n' | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')
  while IFS= read -r value; do
    is_target_device "$value" && return 0
  done <<EOF
$values
EOF
  return 1
}

validate_device() {
  for property in ro.product.device ro.product.name ro.product.vendor.device; do
    is_target_device "$(getprop "$property" 2>/dev/null)" && return 0
  done
  is_target_device_list "$(getprop ro.twrp.target.devices 2>/dev/null)" && return 0
  abort "[ERROR] This firmware is only for: $TARGET_DEVICES"
}

read_bootarg() {
  key="$1"
  for source in /proc/cmdline /proc/bootconfig; do
    [ -r "$source" ] || continue
    value=$(
      sed -e 's/ = /=/g' -e 's/"//g' "$source" |
        tr ' ' '\n' |
        sed -n "s/^${key}=//p" |
        head -n 1
    )
    [ -n "$value" ] && printf '%s\n' "$value" && return 0
  done
  return 1
}

set_slots() {
  active=$(getprop ro.boot.slot_suffix 2>/dev/null)
  [ -n "$active" ] || active=$(getprop ro.boot.slot 2>/dev/null)
  [ -n "$active" ] || active=$(read_bootarg androidboot.slot_suffix)
  [ -n "$active" ] || active=$(read_bootarg androidboot.slot)

  case "$active" in
    a|_a) ACTIVE_SLOT=a; INACTIVE_SLOT=b ;;
    b|_b) ACTIVE_SLOT=b; INACTIVE_SLOT=a ;;
    *) abort "[ERROR] Unable to determine active slot" ;;
  esac
}

check_image() {
  partition="$1"
  expected_hash="$2"
  image="firmware-update/$partition.img"
  output=$(unzip -p "$ZIPFILE" "$image" | sha256sum) ||
    abort "[ERROR] Could not read firmware image: $partition"
  actual_hash=${output%% *}
  [ "$actual_hash" = "$expected_hash" ] ||
    abort "[ERROR] Invalid firmware image: $partition"
}

resolve_partition() {
  partition="$1"
  slot="$2"
  for root in \
    /dev/block/by-name \
    /dev/block/bootdevice/by-name \
    /dev/block/platform/*/by-name
  do
    block="$root/${partition}_${slot}"
    if [ -b "$block" ]; then
      readlink -f "$block"
      return
    fi
  done
  return 1
}

restore_readonly() {
  result=0
  for block in $READONLY_BLOCKS; do
    blockdev --setro "$block" >/dev/null 2>&1 || result=1
  done
  return "$result"
}

cleanup() {
  restore_readonly
  rm -f "$STAGED_IMAGE"
}

preflight_partition() {
  partition="$1"
  image_size="$2"

  for slot in a b; do
    block=$(resolve_partition "$partition" "$slot") ||
      abort "[ERROR] Missing block device: ${partition}_${slot}"
    device_id=$(stat -L -c '%t:%T' "$block" 2>/dev/null) ||
      abort "[ERROR] Could not identify block device: ${partition}_${slot}"
    case " $BLOCK_IDS " in
      *" $device_id "*) abort "[ERROR] Duplicate block device: ${partition}_${slot}" ;;
    esac
    BLOCK_IDS="$BLOCK_IDS $device_id"
    eval "BLOCK_${partition}_${slot}=\$block"
    eval "BLOCK_ID_${partition}_${slot}=\$device_id"

    block_size=$(blockdev --getsize64 "$block" 2>/dev/null) ||
      abort "[ERROR] Could not read partition size: ${partition}_${slot}"
    case "$block_size" in
      ''|*[!0-9]*) abort "[ERROR] Invalid partition size: ${partition}_${slot}" ;;
    esac
    [ "$image_size" -le "$block_size" ] ||
      abort "[ERROR] Image is larger than partition: ${partition}_${slot}"

    read_only=$(blockdev --getro "$block" 2>/dev/null) ||
      abort "[ERROR] Could not read partition state: ${partition}_${slot}"
    case "$read_only" in
      0) ;;
      1)
        blockdev --setrw "$block" >/dev/null 2>&1 ||
          abort "[ERROR] Could not make partition writable: ${partition}_${slot}"
        READONLY_BLOCKS="$READONLY_BLOCKS $block"
        [ "$(blockdev --getro "$block" 2>/dev/null)" = 0 ] ||
          abort "[ERROR] Partition is read-only: ${partition}_${slot}"
        ;;
      *) abort "[ERROR] Invalid partition state: ${partition}_${slot}" ;;
    esac
  done
}

flash_partition() {
  partition="$1"
  image_size="$2"
  expected_hash="$3"
  slot="$4"
  image="firmware-update/$partition.img"
  eval "block=\$BLOCK_${partition}_${slot}"
  eval "expected_device_id=\$BLOCK_ID_${partition}_${slot}"

  unzip -p "$ZIPFILE" "$image" > "$STAGED_IMAGE" ||
    abort "[ERROR] Could not stage firmware image: $partition"
  [ "$(stat -c '%s' "$STAGED_IMAGE" 2>/dev/null)" = "$image_size" ] ||
    abort "[ERROR] Staged image has wrong size: $partition"
  output=$(sha256sum "$STAGED_IMAGE") || abort "[ERROR] Could not hash staged image: $partition"
  actual_hash=${output%% *}
  [ "$actual_hash" = "$expected_hash" ] || abort "[ERROR] Invalid staged image: $partition"

  [ -b "$block" ] || abort "[ERROR] Block device disappeared: ${partition}_${slot}"
  [ "$(stat -L -c '%t:%T' "$block" 2>/dev/null)" = "$expected_device_id" ] ||
    abort "[ERROR] Block device changed: ${partition}_${slot}"
  ui_print "[INFO] Updating ${partition}_${slot}..."
  dd if="$STAGED_IMAGE" of="$block" bs=1048576 conv=fsync ||
    abort "[ERROR] Failed to flash ${partition}_${slot}"
  rm -f "$STAGED_IMAGE"

  output=$(dd if="$block" iflag=count_bytes count="$image_size" 2>/dev/null | sha256sum) ||
    abort "[ERROR] Failed to verify ${partition}_${slot}"
  actual_hash=${output%% *}
  [ "$actual_hash" = "$expected_hash" ] ||
    abort "[ERROR] Verification failed: ${partition}_${slot}"
}

flash_slot() {
  slot="$1"
flash_partition abl 307200 b98346f2cdd697d20b790e93a72dfef940a2219ffae3efc909adec670dccd088 "$slot"
flash_partition aop 335872 c840873a7d159593666655ea079715a8ef3236f3d27cc18e6f42b920dc375421 "$slot"
flash_partition aop_config 28672 212100aab4c659ed59c033bd22f141b130311487f9e7ace61b6e847adead4a2b "$slot"
flash_partition bluetooth 1470464 acc9cf2904280ff5c3ceee09f2f394aeb1e0661fae544de1e098f4e6f002dd64 "$slot"
flash_partition countrycode 1048576 355b102657e92986f185496ee98ce1015f50e1c50f30fff260a23d4831cd2d4f "$slot"
flash_partition cpucp 253952 14fc3684a8fd8698b53f47295c4fedf2219e619d447ee5858c6384cda34bef62 "$slot"
flash_partition cpucp_dtb 16384 ff9092c77bd527ce9e8a84392d372bff045d32c9107718a5c1ee2a024305df7f "$slot"
flash_partition devcfg 61440 08360d73185477491b2f3fb8ef485a15999b8b71fbd9e9c85c8232ce0b7ddd9a "$slot"
flash_partition dsp 67108864 feef74bca3fff9aadc407f9a4aa2a2a6e51998302161059a71675a8896e6ec6e "$slot"
flash_partition featenabler 106496 0c3722c33079278dc7ea50e2be08a10a7e06c5067be13cbb0cba947adfdbe83a "$slot"
flash_partition hyp 1511424 22d3cc8d06acd1afc6f21cf91e8ffacbd0f39d8f6e19c6fab35bf64702ee8ee6 "$slot"
flash_partition imagefv 6299648 de0b78aa60df5c351d30544d45327874fd445a672772be3489499e9c2e55cea7 "$slot"
flash_partition keymaster 442368 447ba86365ee90b86348b0702d45fffde73bbc74b2f86138c4803df9f2d8b1d1 "$slot"
flash_partition modem 156872704 11a1f251c9b174ca16c2a278c62affa85be7dfbd87c8d0a19f45dece6d8fd574 "$slot"
flash_partition modemfirmware 236408832 5e8d086d6871155e3a77f1a0d1ea1a3c8b01b7a90d7ede08f0882b8e2d972300 "$slot"
flash_partition multiimgqti 12288 b268f6b5ae8f8038dca72b6de01282273f81b60e2ed646d063cbf5fa63ecb50b "$slot"
flash_partition qupfw 61440 6e052d85737fc2eec59599d16a23d34934729a632b74d73ea4722eb304711c40 "$slot"
flash_partition shrm 139264 5238a62a4b7b76fcd6d7472d4d64cf76d824043e4dcbf44ee5fbf739060af81b "$slot"
flash_partition spuservice 98304 32146638c4a257db292f16f303f1106585999f9ba16a012f7e78159876ffdf31 "$slot"
flash_partition tz 4141056 a4466bf1323e7fe62ea252a46ce5628040b616cfbd8648f5249853814155c611 "$slot"
flash_partition uefi 2949120 9045ef9ce226f093c58b621414727751515af9591cdae154ee7c90a07d5f551e "$slot"
flash_partition uefisecapp 200704 29fe526f116fb9165776d2dfb40563eb82a128a8b53a03ac3253751d6e84ea20 "$slot"
flash_partition vm-bootsys 20971520 03bced19561f4d369ae2ddc72a013c792c556d93a6dcd7296551c4dd1c118971 "$slot"
flash_partition xbl 1118208 d2eef7a048fb44aee7deac39332862ecf28e3970ab4b550c54259c671a8089b5 "$slot"
flash_partition xbl_config 225280 b1e50d0122d808eb0b2de5fa5d8d45be72fe271fda26e94211855494853ca6e3 "$slot"
flash_partition xbl_ramdump 950272 543038b10f444730d861fa584ad481458edb748fa9a2bd645c0a1fc883eb6447 "$slot"
}

trap cleanup 0
trap 'exit 1' HUP INT TERM

ui_print "####################################################"
ui_print "# Generated by Xiaomi Flashable Firmware Creator"
ui_print "####################################################"
ui_print " "

require_tools
validate_device
set_slots
rm -f "$STAGED_IMAGE"

check_image abl b98346f2cdd697d20b790e93a72dfef940a2219ffae3efc909adec670dccd088
check_image aop c840873a7d159593666655ea079715a8ef3236f3d27cc18e6f42b920dc375421
check_image aop_config 212100aab4c659ed59c033bd22f141b130311487f9e7ace61b6e847adead4a2b
check_image bluetooth acc9cf2904280ff5c3ceee09f2f394aeb1e0661fae544de1e098f4e6f002dd64
check_image countrycode 355b102657e92986f185496ee98ce1015f50e1c50f30fff260a23d4831cd2d4f
check_image cpucp 14fc3684a8fd8698b53f47295c4fedf2219e619d447ee5858c6384cda34bef62
check_image cpucp_dtb ff9092c77bd527ce9e8a84392d372bff045d32c9107718a5c1ee2a024305df7f
check_image devcfg 08360d73185477491b2f3fb8ef485a15999b8b71fbd9e9c85c8232ce0b7ddd9a
check_image dsp feef74bca3fff9aadc407f9a4aa2a2a6e51998302161059a71675a8896e6ec6e
check_image featenabler 0c3722c33079278dc7ea50e2be08a10a7e06c5067be13cbb0cba947adfdbe83a
check_image hyp 22d3cc8d06acd1afc6f21cf91e8ffacbd0f39d8f6e19c6fab35bf64702ee8ee6
check_image imagefv de0b78aa60df5c351d30544d45327874fd445a672772be3489499e9c2e55cea7
check_image keymaster 447ba86365ee90b86348b0702d45fffde73bbc74b2f86138c4803df9f2d8b1d1
check_image modem 11a1f251c9b174ca16c2a278c62affa85be7dfbd87c8d0a19f45dece6d8fd574
check_image modemfirmware 5e8d086d6871155e3a77f1a0d1ea1a3c8b01b7a90d7ede08f0882b8e2d972300
check_image multiimgqti b268f6b5ae8f8038dca72b6de01282273f81b60e2ed646d063cbf5fa63ecb50b
check_image qupfw 6e052d85737fc2eec59599d16a23d34934729a632b74d73ea4722eb304711c40
check_image shrm 5238a62a4b7b76fcd6d7472d4d64cf76d824043e4dcbf44ee5fbf739060af81b
check_image spuservice 32146638c4a257db292f16f303f1106585999f9ba16a012f7e78159876ffdf31
check_image tz a4466bf1323e7fe62ea252a46ce5628040b616cfbd8648f5249853814155c611
check_image uefi 9045ef9ce226f093c58b621414727751515af9591cdae154ee7c90a07d5f551e
check_image uefisecapp 29fe526f116fb9165776d2dfb40563eb82a128a8b53a03ac3253751d6e84ea20
check_image vm-bootsys 03bced19561f4d369ae2ddc72a013c792c556d93a6dcd7296551c4dd1c118971
check_image xbl d2eef7a048fb44aee7deac39332862ecf28e3970ab4b550c54259c671a8089b5
check_image xbl_config b1e50d0122d808eb0b2de5fa5d8d45be72fe271fda26e94211855494853ca6e3
check_image xbl_ramdump 543038b10f444730d861fa584ad481458edb748fa9a2bd645c0a1fc883eb6447

preflight_partition abl 307200
preflight_partition aop 335872
preflight_partition aop_config 28672
preflight_partition bluetooth 1470464
preflight_partition countrycode 1048576
preflight_partition cpucp 253952
preflight_partition cpucp_dtb 16384
preflight_partition devcfg 61440
preflight_partition dsp 67108864
preflight_partition featenabler 106496
preflight_partition hyp 1511424
preflight_partition imagefv 6299648
preflight_partition keymaster 442368
preflight_partition modem 156872704
preflight_partition modemfirmware 236408832
preflight_partition multiimgqti 12288
preflight_partition qupfw 61440
preflight_partition shrm 139264
preflight_partition spuservice 98304
preflight_partition tz 4141056
preflight_partition uefi 2949120
preflight_partition uefisecapp 200704
preflight_partition vm-bootsys 20971520
preflight_partition xbl 1118208
preflight_partition xbl_config 225280
preflight_partition xbl_ramdump 950272

ui_print "[INFO] Updating inactive slot $INACTIVE_SLOT first..."
flash_slot "$INACTIVE_SLOT"
ui_print "[INFO] Updating active slot $ACTIVE_SLOT..."
flash_slot "$ACTIVE_SLOT"

restore_readonly || abort "[ERROR] Failed to restore partition read-only state"
READONLY_BLOCKS=
rm -f "$STAGED_IMAGE"
trap - 0
ui_print "Firmware has been successfully updated!"
