#!/bin/bash
#
# Copyright 2017 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
#
# script to automate the setup steps for 1.6.8 Key increment testing test case
#

# Sample content of conf file
cat <<EOF > /dev/null
# name of the board
BOARD="veyron_jerry"
# Directory where you run cros_sdk
CHROMIUMOS_DIR="/disk2/chromiumos3/"
# location of chromiumos image
IMAGE="/u/dchan/Downloads/jerry/chromiumos_test_image.bin"
# location of the BIOS bin image
BIOS="/u/dchan/Downloads/jerry/image.bin"
LOGS="/u/dchan/Downloads/"
# DUT ip address
DUT_IP="100.96.48.138"

# get_payload uses different name for board.
# BOARD_PAYLOAD="veyron-jerry"
# If you want to bypass the default get_payloads.par
# GETPAYLOAD="/google/data/ro/teams/chromeos-testing/get_payloads.par"
# If you want to bypass the auto generated version string from cros image.
# IMAGE_VER=6812.42.0
# IMAGE_VER_NOSUFFIX=6812.42
# If you want to use a different channel
# CHANNEL=dev
EOF


log()
{
  echo -e $(date '+[%H%M%S-%m%d]') "$*"
}

generate_payload()
{
  #
  # Generate payload in /tmp/key_increment_working_folder
  #
  log
  log ========== generate_payload ==========
  log Copy image $IMAGE to chroot...
  cp $IMAGE $CHROMIUMOS_DIR/chroot/tmp/chromiumos_test_image.bin || exit 1
  log Copy bios $BIOS to chroot...
  cp $BIOS $CHROMIUMOS_DIR/chroot/tmp/image.bin || exit 1

  # Create script to run inside chroot
  log Create /tmp/prepbin.sh inside chroot.
  cat <<EOF > $CHROMIUMOS_DIR/chroot/tmp/prepbin.sh
#!/bin/bash
BOARD=$BOARD
IMAGE=/tmp/chromiumos_test_image.bin
BIOS=/tmp/image.bin
BIOS_STR=$BIOS_STR
BIOS_VER=$BIOS_VER

cd ~/trunk/src/platform/dev/

git log fm_and_key_version_test_prep.sh \
  | grep b681c97226e4e485823e391413208eaa61563291 > /dev/null
if [ \$? -ne 0 ]; then
  echo =======================================================
  echo Your fm_and_key_version_test_prep.sh does not contain the
  echo latest change to work in this script. Please try Tot.
  echo =======================================================
  exit 1
fi
echo Running from chroot \$(pwd)
cmd="./fm_and_key_version_test_prep.sh -b $BOARD_PAYLOAD -i /tmp/chromiumos_test_image.bin -f $BIOS_VER -s /tmp/image.bin -v $BIOS_STR"
\$cmd > /tmp/fm_and_key.out 2>&1

PAYLOAD_BIN="/tmp/key_increment_working_folder/bios/${BOARD_PAYLOAD}_${BIOS_VER}.1_signed.bin"
CK_VER_STR=\$(strings \$PAYLOAD_BIN | grep -m1 '^Google_')
if [ "\$CK_VER_STR" != "\${BIOS_STR}.test1" ]; then
  echo resulting test perp version mismatch
  echo \$PAYLOAD_BIN version is \$CK_VER_STR, expects \${BIOS_STR}.test1
  echo Please investigate.
  exit 1
fi
EOF
  chmod ugo+x $CHROMIUMOS_DIR/chroot/tmp/prepbin.sh

  if [ -d $CHROMIUMOS_DIR/chroot/tmp/key_increment_working_folder ]; then
    rm -rf $CHROMIUMOS_DIR/chroot/tmp/key_increment_working_folder
  fi
  log Launching cros_sdk /tmp/prepbin.sh, this will take 9-10 minutes.
  log For progress: tail -f $CHROMIUMOS_DIR/chroot/tmp/fm_and_key.out
  (
  date
  cd $CHROMIUMOS_DIR
  cros_sdk -- bash /tmp/prepbin.sh < /dev/null
  date
  )
  log
  log "Check the output:"
  log "  $CHROMIUMOS_DIR/chroot/tmp/fm_and_key.out"
  log "You can ignore the error near the end"
  log "Rename testimage to $CHANNEL."
  for i in  $CHROMIUMOS_DIR/chroot/tmp/key_increment_working_folder/payloads/*.signed
  do
    mv $i $(echo $i|sed 's/_testimage-/_'$CHANNEL'-/g')
  done
}


upload_payload()
{
  #
  # Upload the payload generated by generate_payload
  #
  log
  log ========== upload_payload ==========
  gsutil_location="gsutil"
  log Looking for gsutil
  while true; do
    which gsutil
    if [ $? -eq 0 ]; then
      break
    fi
    log Failed to locate $gsutil_location.
    log Please enter the full path of gsutil:
    read gsutil_location
    if $gsutil_location -v; then
      break
    fi
  done
  log Found gsutil $gsutil_location

  cmd="gsutil update"
  log "Run $cmd to fetch the latest gsutil, answer y if prompted."
  $cmd

  # upload
  cmd="$gsutil_location -m cp -R $CHROMIUMOS_DIR/chroot/tmp/key_increment_working_folder/payloads/* gs://chromeos-throw-away-bucket/CrOSPayloads/$STORE_DIR/"
  log "Uploading payload $cmd"
  $cmd

  log
  log Your files has been uploaded to https://pantheon.corp.google.com/storage/browser/chromeos-throw-away-bucket/CrOSPayloads/$STORE_DIR/
  log

  # enable shared publicly
  cmd="$gsutil_location -m acl ch -u AllUsers:R gs://chromeos-throw-away-bucket/CrOSPayloads/$STORE_DIR/*"
  log "Make file shared publicly $cmd"
  $cmd

  # check shared publicly
  err=0
  log Check perm is set correctly
  for f in $CHROMIUMOS_DIR/chroot/tmp/key_increment_working_folder/payloads/*; do
    cmd="$gsutil_location acl get gs://chromeos-throw-away-bucket/CrOSPayloads/$STORE_DIR/$(basename $f)"
    $cmd | grep allUsers > /dev/null
    if [ $? -ne 0 ] ; then
      log Failed to enable Shared Publicly for $f
      err=$((err+1))
    fi
  done

  if [ $err -gt  0 ]; then
    log There appear to be some problem with setting the Shared Publicly flag, please check error above.
    log Press enter when ready to continue...
    read ans
  fi
}


generate_config()
{
  #
  # Generate the configuration for omaha dev server
  #
  log
  log ========== generate_config ==========
  if [ -z "$GETPAYLOAD" ]; then
    NONCONF="/home/build/nonconf/google3/"
    GET_PAYLOAD="/google/data/ro/teams/chromeos-testing/qa_au.par"
  else
    GET_PAYLOAD=${GETPAYLOAD}
  fi

  # TODO: we assume cros bin ends with .0 and the continue value is 1-5
  # if you OS ends with other digit,
  # make sure the suffix is of increasing value.
  # remove the old one otherwise payloads.py will genreate duplicate Rules.
  rm $FINAL_CONFIG 2> /dev/null
  rm autoupdate-ascii-chromeos-*-${IMAGE_VER_NOSUFFIX}.[1-5].config \
    2> /dev/null
  log Done with cleanup, your current directory listing:
  ls
  COMMON_DATA="commondatastorage.googleapis.com"
  COMMON_PATH="/chromeos-throw-away-bucket/CrOSPayloads/$STORE_DIR"
  PAYLOAD_FOLDER=/tmp/key_increment_working_folder/payloads
  for i in {1..5}; do
    set -x
    $GET_PAYLOAD \
      -p $BOARD_PAYLOAD  \
      -b ${IMAGE_VER_NOSUFFIX}.$i \
      -c $CHANNEL-channel \
      -u "https://${COMMON_DATA}${COMMON_PATH}" \
      --keyinctest \
      --folder "${CHROMIUMOS_DIR}/chroot${PAYLOAD_FOLDER}" \
      --board test \
      $GETPAYLOAD_ARGS
    set +x
  done > /tmp/get_payloads.out 2>&1

  log If there are problem, check /tmp/get_payloads.out
  # generate a single config
  (
  pat='^AppId|^ConfigName|^IsPublic|^EnableTargetVersionCheck';
  egrep "$pat" \
    autoupdate-ascii-chromeos-${BOARD_PAYLOAD}-${IMAGE_VER_NOSUFFIX}.1.config
  egrep -hv "$pat" \
    autoupdate-ascii-chromeos-${BOARD_PAYLOAD}-${IMAGE_VER_NOSUFFIX}.[1-5].config \
    | awk '{ n=match($0, "-[0-9]+\\.[0-9]+\\.[0-9])");
             if(n) {lstv=v; v=substr($0,n+1,RLENGTH-2);}
           }
           { if(lstv) {
               if(match($0, "0[\\.0]+\\.0-")){
                 sub("0[.0]+.0", lstv, $0); lstv="";
               }
             }
             print
           }'
  ) > $FINAL_CONFIG

  # Copy autoconfig to google data for easy copy and paste
  (cd $GOOGLEDATA)  # to automount /google/data
  cp $FINAL_CONFIG $GOOGLEDATA || log Failed to copy $FINAL_CONFIG $GOOGLEDATA

  log Version updated:
  grep ' Version:' $FINAL_CONFIG
  log Check the content of $FINAL_CONFIG
  log Press ENTER when done...
  read ans
}


upload_config()
{
  #
  # Upload the config from genreate_config() to omaha dev server
  #
  log
  log ========== upload_config ==========
  if [ -z "$DUT_IP" ]; then
    ans="s"
  else
    cat <<EOF

    Running upload_config

    1 Device in normal mode
    2 Connected to corp network with IP $DUT_IP
    3 chromeOS version is $IMAGE_VER
    4 Enable hardware and software write protect
    Press ENTER when done (s to skip)...
EOF
    read ans
  fi

  > ~/.ssh/known_hosts

  s="CHROMEOS_AUSERVER=https://omaha.sandbox.google.com/service/update2"
  if [ "$ans" == "s" ]; then
    log Skip DUI setup.
    log You will need to execute the following on DUT:
    log "echo $s > /mnt/stateful_partition/etc/lsb-release"
    log
  else
    set -x
    ssh root@$DUT_IP "echo $s > /mnt/stateful_partition/etc/lsb-release" \
      || exit 1
  fi
  set +x

  # In case copy to x20 failed, tell user to copy from local file.
  if [ -f "$GOOGLEDATA/FINAL_CONFIG" ]; then
    final_config="https://x20web.corp.google.com/~$USER/$FINAL_CONFIG"
  else
    final_config=$FINAL_CONFIG
  fi

  cat <<EOF
  Detail in http://omahauploadconfig.  Steps are:
  1. Send email (cc wireless-hw-testing@ for storm)
     To: chromeos-test@google.com
     Subject: Key increment test for ${BOARD}
     Content:
     Starting ${BOARD} Key Increment testing and we are using
     {Insert https://omahaconsole-dev.corp.google.com/ config URL}
     Reply to this thread if you need to reconfigure or update.
     Thank you.
  2. Open https://omahaconsole-dev.corp.google.com/
  3. Click on the row for your board.
  4. Click "Create New" button in the Proposed Version box
  5. Paste content from $final_config
  6. Click Create button.
  7. Click VERIFY CONFIG and check that the web page returns OK.
  8. Check email that no one object on omaha update.
  9. Click "Push to Live" button in Proposed Version box.

  Press ENTER when done ...
EOF
  read ans
}


exit_on_error()
{
  echo $(date '+[%H%M%S-%m%d]') ERROR:- "$*"
  exit 1
}


get_logs()
{
  #
  # Get the logs from DUT on each auto update sucess/failure.
  #

  scp root@$DUT_IP:\{/var/log/\{update_engine.log,debug_vboot_noisy.log\},/etc/lsb-release\} $(pwd)/key_increment_test/AU$1/
  if [ $? -ne 0 ]; then
    exit_on_error "Check network and DUT status"
  fi
  echo $2
}


debug_log_verification()
{
  #
  # GPT_INDEX and PART_NUM verification from /var/log/debug_vboot_noisy.log file
  #

  GPT_INDEX=$(ssh root@$DUT_IP grep -m1 "GPT\ index" /var/log/debug_vboot_noisy.log)
  GPT_INDEX=$(cut -d= -f2- <<<$GPT_INDEX)

  PART_NUM=$(ssh root@$DUT_IP grep -m1 "KERN-B" /var/log/debug_vboot_noisy.log)
  PART_NUM=$(cut -d" " -f 3 <<<$PART_NUM)

  if [ $GPT_INDEX -ne 4 ] || [ $PART_NUM -ne 4 ]; then
    get_logs 1 "debug_log_verification failed on AU1"
    exit_on_error "GPT_INDEX= $GPT_INDEX, PART_NUM= $PART_NUM"
  fi
}


crossystem_verification()
{

  #
  # CHROMEOS_RELEASE_VERSION update verification from /etc/lsb-release
  # FWB_TRIES and FWID parameters verification from crossystem
  #

  log ----inside crossystem_verification $1

#  local au_iteration=$1

  CHROMEOS_RELEASE_VERSION=$(ssh root@$DUT_IP grep "CHROMEOS_RELEASE_VERSION" /etc/lsb-release)
  CHROMEOS_RELEASE_VERSION=$(cut -d= -f2- <<<$CHROMEOS_RELEASE_VERSION)
  if [ "$CHROMEOS_RELEASE_VERSION" != "${IMAGE_VER_NOSUFFIX}.$1" ]; then
    get_logs $auto_update "CHROMEOS_RELEASE_VERSION match Failed on AU$1"
    exit_on_error "CHROMEOS_RELEASE_VERSION=$CHROMEOS_RELEASE_VERSION"
  fi

  FWB_TRIES=$(ssh root@$DUT_IP crossystem fwb_tries)
  if [ $FWB_TRIES -eq 0 ]; then
    get_logs $auto_update " FWB_TRIES is '0' on AU$1"
    exit_on_error "FWB_TRIES=$FWB_TRIES"
  fi

  FWID=$(ssh root@$DUT_IP crossystem fwid)
  if [ "$FWID" != "${BIOS_STR}.test$1" ]; then
    get_logs $auto_update " FWID match Failed on AU$1"
    exit_on_error "FWID=$FWID"
  fi

}


verify_keyincrement()
{
  #
  # Verify key increment test on each auto update.
  #

  mkdir -p $(pwd)/key_increment_test/AU{1..5}

  for auto_update in {1..5}; do
    log ----inside verify_keyincrement $auto_update

    ssh root@$DUT_IP update_engine_client  --omaha_url="https://omaha.sandbox.google.com/service/update2" --check_for_update --block_until_reboot_is_needed

    #ssh root@$DUT_IP update_engine_client  --check_for_update --block_until_reboot_is_needed

    if [ $? -ne 0 ]; then
      get_logs $auto_update "update_engine_client Failed on AU$auto_update"
      exit 1
    fi

    ssh root@$DUT_IP reboot

    sleep 45
    crossystem_verification $auto_update
    sleep 60

    ssh root@$DUT_IP "grep completed /var/log/update_engine.log"

    if [ $? -ne 0 ]; then
      log Autoupdate Verification Failed on AU$auto_update after reboot
      exit 1
    fi

    if [ $auto_update -eq 1 ]; then
      debug_log_verification
    fi

    ssh root@$DUT_IP "reboot"

    sleep 90
    get_logs $auto_update "Autoupdate $auto_update completed"

  done

  # Final verification
  # tpm_fwver and tpm_kernver verification.

  TPM_FWVER=$(ssh root@$DUT_IP crossystem tpm_fwver)
  echo "TPM_FWVER=$TPM_FWVER"
  TPM_KERNVER=$(ssh root@$DUT_IP crossystem tpm_kernver)
  echo "TPM_KERNVER=$TPM_KERNVER"

  if [ "$TPM_FWVER" != "0x00030003" ] || [ "$TPM_KERNVER" != "0x00030003" ]; then
    get_logs 0 "TPM_FWVER=$TPM_FWVER and TPM_KERNVER=$TPM_KERNVER"
    exit_on_error "TPM_FWVER=$TPM_FWVER and TPM_KERNVER=$TPM_KERNVER"
  fi

  log "Key Increment Test PASS\n\n"

  tar -zcf key_increment_test.tar.gz key_increment_test
  rm -rf key_increment_test
}


recover_device()
{

  cd $CHROMIUMOS_DIR

  log "Copy firmware to DUT\n\n"
  scp $CHROMIUMOS_DIR/chroot/tmp/image.bin root@$DUT_IP:/tmp/

  log "Flashing firmware\n\n"

  ssh root@$DUT_IP flashrom -p host -w /tmp/image.bin -i RW_SECTION_A -i RW_SECTION_B

  sleep 20
  log "Flashing testimage\n\n"
  cros flash $DUT_IP $IMAGE
  ssh root@$DUT_IP reboot

  log "Clear TPM\n\n"

}


#----------------------------------------------------------------------
# MAIN
#----------------------------------------------------------------------
CONFIG=$1
FUNC=$2
if [ "$CONFIG" = "" ]; then
  log You must specify config file
  exit 1
fi

if [ "$FUNC" = "" ]; then
  cat <<EOF
You must provide a method to execute.
Available functions:
  generate_payload
  upload_payload
  generate_config
  upload_config
  verify_keyincrement
  recover_device
EOF
  exit 1
fi
. $CONFIG

if [ -z "$CHANNEL" ]; then
  CHANNEL=dev
fi


#-----------------------------------------------------------
# Let's get the sudo and prodaccess password out of the way
#-----------------------------------------------------------
log type in sudo password if prompt
sudo pwd
log prod access if prompt
prodcertstatus > /dev/null || prodaccess

#-----------------------------------------------------------
# Generated values
#-----------------------------------------------------------
# Get firmware version
BIOS_STR=$(strings $BIOS | grep -m1 '^Google_')
BIOS_VER=$(cut -d. -f2- <<<$BIOS_STR)
log BIOS version: $BIOS_VER

# Get OS version from chromiumos_test_image.bin
if [ -z "$IMAGE_VER" ]; then
  log Extract chromeOS version and set IMAGE_VER
  s='^CHROMEOS_RELEASE_VERSION='
  IMAGE_VER=$(strings $IMAGE|grep -m1 "$s"|cut -d= -f2)
else
  log Use user defined IMAGE_VER
fi
# IMAGER_VER should looks Ex: 7179.0.0
log Chrome version: $IMAGE_VER

if [ -z "$IMAGE_VER_NOSUFFIX" ]; then
  # Remove the last digit Ex: 7179.0
  log Set chromeOS version prefix and set IMAGE_VER_NOSUFFIX
  s='^CHROMEOS_RELEASE_VERSION='
  IMAGE_VER_NOSUFFIX=$(awk -F. 'sub(FS $NF"$", x)' <<<$IMAGE_VER)
else
  log Use user defined IMAGE_VER_NOSUFFIX
fi
log Chrome version without suffix: $IMAGE_VER_NOSUFFIX

if [ -z "$BOARD_PAYLOAD" ]; then
  BOARD_PAYLOAD="$BOARD"
fi
log Board name: $BOARD
log Board prepartion name: $BOARD_PAYLOAD
log Will use chroot: $CHROMIUMOS_DIR

# directory name in storage
STORE_DIR=${BOARD}_keyinc_fw$(sed 's/\./_/g' <<<$BIOS_VER)_$USER
FINAL_CONFIG="autoupdate-ascii-chromeos-$BOARD.config"
GOOGLEDATA="/google/data/rw/users/$(cut -c1-2 <<<$USER)/$USER/www"

if [ "$FUNC" = "ALL" ]; then
  generate_payload
  upload_payload
  generate_config
  upload_config
  verify_keyincrement
  recover_device
  exit 0
fi

$FUNC
exit 0