#!/bin/bash

#
# Copyright 2022-2023 Nestybox, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    https:#www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

# Usage:
#        sysbox-docker-cp [OPTIONS] CONTAINER:SRC_PATH DEST_PATH|-
#        sudo sysbox-docker-cp [OPTIONS] SRC_PATH|- CONTAINER:DEST_PATH
#
# Notes:
#
# * When copying from the container to the host, use
#   "sysbox-docker-cp" not "sudo sysbox-docker-cp" as otherwise the
#   file will be chowned to "root:root" on the host.
#
# * When copying from the host to the container, use
#   "sudo sysbox-docker-cp", as this script needs root privileges on the
#   host in order to access the container's root filesystem (typically
#   under /var/lib/docker/).
#
# * Options are identical to "docker cp"; refer to the "docker cp --help"
#   for further info.
#
# "sysbox-docker-cp" implements "docker cp" for Sysbox containers. It's needed
# because "docker cp" does not work well in containers with the user namespace
# (unless Docker is configured with userns-remap, which is uncommon).
#
# The reason "docker cp" does not work well with Sysbox containers is
# that it's not aware that the Sysbox container uses the
# user-namespace, so it performs the copy of the file but does not
# "chown" it according to the user-namespace mapping. Thus, the file
# shows up as "nobody:nogroup" inside the container, which is not as
# intended. Furthermore, the file can't be chowned from within the
# container (permission denied).
#
# As a work-around, this tool (sysbox-docker-cp) peforms the copy and chown. It
# simply calls "docker cp" and then performs the chown of the file.
#


# Globals

# Note: in the future, get the sysbox data root from sysbox itself.
SYSBOX_DATA_ROOT="/var/lib/sysbox"
ARCHIVE=""
FOLLOW_LINK=""

function show_usage() {
  printf "This tool is a wrapper for \"docker cp\" for containers created with Docker + Sysbox.\n"
  printf "It's needed to copy with correct user:group ID's within the Sysbox container's \n"
  printf "user-namespace (\"docker cp\" does not do this). Usage is similar to \"docker cp\".\n"
  printf "\n"
  printf "To copy from container to the host:\n"
  printf "sysbox-docker-cp [OPTIONS] CONTAINER:SRC_PATH DEST_PATH|-\n"
  printf "\n"
  printf "To copy from the host to the container, you must use sudo:\n"
  printf "sudo sysbox-docker-cp [OPTIONS] SRC_PATH|- CONTAINER:DEST_PATH\n"
  printf "\n"
  printf "Options are identical to \"docker cp\"; refer to the \"docker cp --help\" for further info.\n"
  printf "\n"
}

function parse_opt() {
  options=$(getopt -o aLh -l archive,follow-link,help -- "$@")

  eval set -- "$options"

  while true; do
	  case "$1" in
	    -h | --help)
		    show_usage
		    exit 1
		    ;;
	    -a | --archive)
		    ARCHIVE="-a"
		    ;;
	    -L | --follow-link)
		    FOLLOW_LINK="-L"
		    ;;
	    --)
		    shift
		    break
		    ;;
	    -*)
		    show_usage
		    exit 1
		    ;;
	    *)
		    show_usage
		    exit 1
		    ;;
	  esac
	  shift
  done

  if [[ $# != 2 ]]; then
	  show_usage
	  exit 1
  fi

  ARG1=$1
  ARG2=$2
}

function parse_args() {
  if [[ "$ARG1" == *":"* ]]; then
	  COPY_TO_CONTAINER=false
	  CONT_NAME=$(echo $ARG1 | cut -d ":" -f1)
	  SRC_PATH=$(echo $ARG1 | cut -d ":" -f2)
	  DST_PATH=$ARG2
  elif [[ "$ARG2" == *":"* ]]; then
	  COPY_TO_CONTAINER=true
	  CONT_NAME=$(echo $ARG2 | cut -d ":" -f1)
	  DST_PATH=$(echo $ARG2 | cut -d ":" -f2)
	  SRC_PATH=$ARG1
  else
	  printf "\nError: incorrect argument; was the container specified?\n"
	  exit 1
  fi
}

function docker_cp() {
  docker cp ${ARCHIVE} ${FOLLOW_LINK} ${ARG1} ${ARG2}
  if [[ $? != 0 ]]; then
	  exit 1
  fi
}

function is_sysbox_container() {
  local runtime=$(docker container inspect --format '{{.HostConfig.Runtime}}' $CONT_NAME)
  [[ "$runtime" == "sysbox-runc" ]]
}

function sysbox_rootfs_cloned() {
  stat ${SYSBOX_DATA_ROOT}/rootfs/${CONT_ID} > /dev/null 2>&1
}

function sysbox_rootfs_idmapped() {
  kernel_version=$(uname -r | cut -d- -f1)
  # overlayfs on idmapped mounts is supported in kernel 5.19 or above
  if [ "$(echo -e "$kernel_version\n5.19" | sort -V | head -n1)" = "5.19" ]; then
	  return 0
  else
	  return 1
  fi
}

function sysbox_rootfs_shiftfs() {
  lsmod | grep shiftfs > /dev/null 2>&1
}

function get_container_uid() {
  cat /proc/${CONT_INIT_PID}/uid_map | awk '{print $1}'

}

function chown_recursive() {
  local path="$1"
  local uid_offset="$2"
  local gid_offset="$3"

  # Get current UID and GID
  current_uid=$(stat -c "%u" "$path")
  current_gid=$(stat -c "%g" "$path")

  # Calculate new UID and GID
  new_uid=$((current_uid + uid_offset))
  new_gid=$((current_gid + gid_offset))

  # Change ownership of the current path
  chown "$new_uid:$new_gid" "$path"

  # If it's a directory, recurse into its contents
  if [[ -d "$path" ]]; then
    for item in "$path"/*; do
      # Skip if the directory is empty
      [[ -e "$item" ]] || continue

      # Recursively process the item
      chown_recursive "$item" $CONT_HOST_UID $CONT_HOST_GID
    done
  fi
}

function chown_container_file() {
  local path="${ROOTFS}/${DST_PATH}"

  # if the DST_PATH ends with ".", then append to it the copied file/dir
  if [[ "$DST_PATH" =~ \.$ ]]; then
    local src_base=$(basename ${SRC_PATH})
    path="${ROOTFS}/${DST_PATH}/${src_base}"
  fi

  chown_recursive $path $CONT_HOST_UID $CONT_HOST_GID
}

function chown_local_file() {
  local local_path=${DST_PATH}

  # if the DST_PATH ends with ".", then append to it the copied file/dir
  if [[ "$DST_PATH" =~ \.$ ]]; then
    local src_base=$(basename ${SRC_PATH})
    local_path="${DST_PATH}/${src_base}"
  fi

  real_user_uid=$(id -u)
  real_user_gid=$(id -g)

  # Chown the copied files (recursively) to the user's ID.
  ret=$(chown -R ${real_user_uid}:${real_user_gid} $local_path)
  if [[ $? != 0 ]]; then
	  exit 1
  fi
}

function docker_using_userns_remap() {
  local usernsMode=$(docker inspect --format='{{.HostConfig.UsernsMode}}' $CONT_NAME)
  [[ "$usernsMode" = "userns" ]]
}

function get_container_info() {
  CONT_ID=$(docker inspect --format '{{.Id}}' $CONT_NAME)
  CONT_INIT_PID=$(docker inspect -f '{{.State.Pid}}' $CONT_ID)
  CONT_HOST_UID=$(cat /proc/${CONT_INIT_PID}/uid_map | awk '{print $2}')
  CONT_HOST_GID=$(cat /proc/${CONT_INIT_PID}/gid_map | awk '{print $2}')

  # order is important here
  if docker_using_userns_remap; then
	  SYSBOX_ROOTFS_TYPE="unmapped"
	  ROOTFS=$(docker container inspect --format '{{.GraphDriver.Data.MergedDir}}' $CONT_NAME)
  elif sysbox_rootfs_cloned; then
	  SYSBOX_ROOTFS_TYPE="cloned"
	  ROOTFS="${SYSBOX_DATA_ROOT}/rootfs/${CONT_ID}/overlay2/merged/"
  elif sysbox_rootfs_idmapped; then
	  SYSBOX_ROOTFS_TYPE="idmapped"
	  ROOTFS=$(docker container inspect --format '{{.GraphDriver.Data.UpperDir}}' $CONT_NAME)
  elif sysbox_rootfs_shiftfs; then
	  SYSBOX_ROOTFS_TYPE="shiftfs"
	  ROOTFS=$(docker container inspect --format '{{.GraphDriver.Data.MergedDir}}' $CONT_NAME)
  else
	  SYSBOX_ROOTFS_TYPE="unmapped"
	  ROOTFS=$(docker container inspect --format '{{.GraphDriver.Data.MergedDir}}' $CONT_NAME)
  fi
}

function main() {
  parse_opt "$@"
  parse_args

  get_container_info

  if ! is_sysbox_container; then
	  printf "Error: ${CONT_NAME} is not a Sysbox container.\n"
	  exit 1
  fi

  if [[ $COPY_TO_CONTAINER == "true" ]]; then
    stat ${ROOTFS} > /dev/null 2>&1
    if [ $? -ne 0 ]; then
      printf "Error: must have root access (e.g., sudo) to copy to container.\n"
      exit 1
    fi
  fi


  docker_cp

  # If sysbox is using rootfs cloning or ID-mapped mounts, we
  # need to chown the copied file.
  if [[ $SYSBOX_ROOTFS_TYPE == "cloned" ]] || [[ $SYSBOX_ROOTFS_TYPE == "idmapped" ]]; then
	  if [[ $COPY_TO_CONTAINER == true ]]; then
	    chown_container_file
	  else
	    chown_local_file
	  fi
  fi

  exit 0
}

main "$@"
