#!/bin/bash
#------------------------------------------------------------
# - Save docker images to disk or load images from disk
# - Pull docker images from 'gcr.io|k8s.gcr.io|quay.io'
#
# @author:  gjmzj
# @usage:   ./imgutils
# @ref:     https://github.com/easzlab/kubeasz/tree/master/tools/imgutils

set -o nounset
set -o errexit
#set -o xtrace

function usage() {
  echo -e "\033[33mUsage:\033[0m imgutils [options] [args]"
  cat <<EOF
  option: -{LPSdr}
    -L         to load images from the disk (default '/tmp/docker/images')
    -P         to pull images from CN mirrors of 'docker.io|gcr.io|k8s.gcr.io|quay.io'
    -S         to save local images to the disk (default '/tmp/docker/images')
    -d <dir>   change image directory (default '/tmp/docker/images')
    -r <str>   limit to images who's repo have a pattern <str>

Example:
$ ./imgutils -S -d /opt/docker/images -r gcr.io
  - to save all local images of repo 'gcr.io' into '/opt/docker/images/'
$ ./imgutils -P k8s.gcr.io/addon-resizer:1.8.3
  - to pull images from k8s.gcr.io
EOF
}

function logger() {
  TIMESTAMP=$(date +'%Y-%m-%d %H:%M:%S')
  case "$1" in
    debug)
      echo -e "$TIMESTAMP \033[36mDEBUG\033[0m $2"
      ;;
    info)
      echo -e "$TIMESTAMP \033[32mINFO\033[0m $2"
      ;;
    warn)
      echo -e "$TIMESTAMP \033[33mWARN\033[0m $2"
      ;;
    error)
      echo -e "$TIMESTAMP \033[31mERROR\033[0m $2"
      ;;
    *)
      ;;
  esac
}


function save_to_disk() {
  if [[ "$REPO_PATTERN" == "" ]];then
    DOCKER_IMAGES=$(docker images|grep -v "^REPOSITORY"|awk '{print $1":"$2}')
  else
    DOCKER_IMAGES=$(docker images|grep "$REPO_PATTERN"|awk '{print $1":"$2}')
  fi
  
  # trim the last '/' in "$IMAGE_DIR"
  IMAGE_DIR=$(echo "$IMAGE_DIR"|sed 's/\/$//g')
  mkdir -p "$IMAGE_DIR"

  for image in ${DOCKER_IMAGES};do
    name=$(echo "$image"|sed 's/\//_/g'|sed 's/:/_/g')
    name="$IMAGE_DIR/$name.tar"
    logger info "packing $name..."
    docker save -o "$name" "$image"
  done
}

function load_from_disk() {
  [[ -d "$IMAGE_DIR" ]] || { logger error "Invald directory $IMAGE_DIR"; exit 1; }
  # trim the last '/' in "$IMAGE_DIR"
  IMAGE_DIR=$(echo "$IMAGE_DIR"|sed 's/\/$//g')
  
  if [[ "$REPO_PATTERN" == "" ]];then
    DOCKER_PACKS="$IMAGE_DIR/*.tar"
  else
    DOCKER_PACKS="$IMAGE_DIR/*$REPO_PATTERN*.tar"
  fi
 
  for pack in ${DOCKER_PACKS};do
    logger info "loading $pack..."
    docker load -i "$pack"
  done
}

function pull_any_image() {
  F=$(echo "$1"|awk -F/ '{print NF}')
  case "$F" in
    1)
      HUB="docker.mirrors.ustc.edu.cn"
      GRP="library"
      IMG="$1"
      ;;
    2)
      F1=$(echo "$1"|awk -F/ '{print $1}')
      F2=$(echo "$1"|awk -F/ '{print $2}')
      if [[ "$F1" == k8s.gcr.io ]];then
        HUB="gcr.mirrors.ustc.edu.cn"
        GRP="google-containers"
        IMG="$F2"
      else
        HUB="docker.mirrors.ustc.edu.cn"
        GRP="$F1"
        IMG="$F2"
      fi
      ;;
    3)
      F1=$(echo "$1"|awk -F/ '{print $1}')
      F2=$(echo "$1"|awk -F/ '{print $2}')
      F3=$(echo "$1"|awk -F/ '{print $3}')
      if [[ "$F1" == quay.io ]];then
        HUB="quay.mirrors.ustc.edu.cn"
      elif [[ "$F1" == gcr.io ]];then
        HUB="gcr.mirrors.ustc.edu.cn"
      else
        logger warn "no mirror image for $1"
        HUB="$F1"
      fi
      GRP="$F2"
      IMG="$F3"
      ;;
    *)
      logger error "Invald image $1"
      exit 1
      ;;
  esac
  logger info "Using mirror $HUB/$GRP/$IMG"
  docker pull "$HUB/$GRP/$IMG" && \
  docker tag "$HUB/$GRP/$IMG" "$1" && \
  docker rmi "$HUB/$GRP/$IMG"
}

function main() {
  # check if use bash shell
  readlink /proc/$$/exe|grep -q "dash" && { logger error "you should use bash shell, not sh"; exit 1; }
  # check if use with root
  [[ "$EUID" -ne 0 ]] && { logger error "you should run this script as root"; exit 1; }
  # check num of args
  [[ "$#" -eq 0 ]] && { usage >&2; exit 1; }
  
  export IMAGE_DIR="/tmp/docker/images"
  export REPO_PATTERN=""
  export ACTION=""

  while getopts "LP:Sd:r:" OPTION; do
      case "$OPTION" in
        L)
          ACTION="load_from_disk"
          ;;
        P)
          ACTION="pull_any_image $OPTARG"
          ;;
        S)
          ACTION="save_to_disk"
          ;;
        d)
          export IMAGE_DIR="$OPTARG"
          ;;
        r)
          export REPO_PATTERN="$OPTARG"
          ;;
        ?)
          usage
          exit 1
          ;;
      esac
  done
  
  [[ "$ACTION" == "" ]] && { logger error "illegal option"; usage; exit 1; }
  
  # excute cmd "$ACTION" 
  logger debug "Action begin: $ACTION"
  ${ACTION} || { logger error "Action failed: $ACTION"; return 1; }
  logger debug "Action successed: $ACTION"
}

main "$@"