#!/usr/bin/env bash
#
# This program provides a convenient utility for "building" a compatible 
# application image from pre-existing application image formats,
# such as Docker images or QEMU VHDs.
#

# Initialize Configuration default values
# Can only be string value "true" or "false"
CHECK_SYMBOL=false
USE_MOUNT=false
VERBOSE=false

function print_usage {
    echo "$(basename $0): a useful shell script for converting Docker images into Mystikos cpio archives"
    echo
    echo "Usage: $(basename $0) [-o file] [-i image] [-e extras] [-v] [[-d] file]"
    echo
    echo "  -d file     Build docker image from the named file"
    echo "  -i image    Specify an existing Docker image name to build from"
    echo "  -f format   Specify the output format. Options: dir, cpio. Default: 'dir'"
    echo "  -e extras   Specify the extra options while building the container image"
    echo "  -o      Specify the name of the output file or directory"
    echo "          If not specified, defaults to 'appdir'"
    echo "  -m      Export docker image by mount. This is faster but requires sudo."
    echo "  -v      Verbose"
    echo "  -p      Print unsupported glibc symbol present in application"
    echo "  -h      Print this help message"
    echo
    echo
    echo "Sample Usage:"
    echo
    echo "  Build a Dockerfile without '-d' or '-o' option:"
    echo "    $(basename $0) Dockerfile"
    echo
    echo "  Build Dockerfile with irregular file name"
    echo "    $(basename $0) Dockerfile.release"
    echo
    echo "  Build existing image from remote registry"
    echo "    $(basename $0) -i hello-world"
    echo
    echo "  Build specified Docker image with extra options in verbose"
    echo "    $(basename $0) -d Dockerfile.debug -v -e '--build-arg USERNAME=username --build-arg PASSWORD=password'"
    echo
}

# Please add appropriate cleanup actions here for any new resources created in this script
function cleanup()
{
    # this is a catchall, ignore errors
    set +e
    [ -f $TEMP_IMAGE_IIDFILE ] && rm $TEMP_IMAGE_IIDFILE
    [ -d $EXTRACTED_ROOTFS_DIR ] && rm -rf $EXTRACTED_ROOTFS_DIR
    if [ "$USE_MOUNT" = "true" ]; then
        if [ -d $MNT_PNT_DIR ]; then
            sudo umount $MNT_PNT_DIR
            rm -rf $MNT_PNT_DIR
        fi
    else
        [ -f $EXPORTED_IMG_TAR ] && rm $EXPORTED_IMG_TAR
    fi

}

trap cleanup EXIT

while getopts 'hpvmd:i:f:o:e:' OPTION; do
    case "$OPTION" in
        d )
            DOCKER_FILE=${OPTARG:-Dockerfile}
            ;;
        f )
            OUTPUT_FORMAT="$OPTARG"
            ;;
        h )
            print_usage;
            exit 0
            ;;
        i )
            IMAGE_NAME="$OPTARG"
            ;;
        o )
            OUTPUT_NAME="$OPTARG"
            ;;
        e )
            EXTRA_ARGS="$OPTARG"
            ;;
        v )
            VERBOSE=true;
            set -x
            ;;
        m )
            USE_MOUNT=true;
            ;;
        p )
            CHECK_SYMBOL=true;
            ;;
        * )
            print_usage;
            exit 1
            ;;
    esac
done

### defaults

if [ $# -lt 1 ]; then
        echo "Dockerfile or ImageName has not been specified"
        exit 1
fi

# if neither '-i' nor '-d' option is specified, assume the last argument is a Dockerfile name. 
# ImageName/URL must be specified by '-i' option. The default argument can only be a Dockerfile
if [ -z "$DOCKER_FILE" -a -z "$IMAGE_NAME" ]; then
    DOCKER_FILE=${@: -1}
fi

OUTPUT_NAME=${OUTPUT_NAME:-appdir}
DELETE_OUTPUT=${DELETE_OUTPUT:-false}
OUTPUT_FORMAT=${OUTPUT_FORMAT:-dir}
TEMP_IMAGE_IIDFILE=$(mktemp -t myst-iid.XXXXXX)
NETWORK_NAME=$(basename ${PWD})-$(head /dev/urandom | tr -dc 'a-z0-9' | head -c 10)
DOCKER_RETRIES=3

### end defaults

set -e

if [ ! -z "$DOCKER_FILE" ]; then
    USE_DOCKERFILE=true
else
    USE_DOCKERFILE=false
fi

case "$OUTPUT_FORMAT" in
    dir )
        ;;
    cpio )
        ;;
    ext2 )
        echo "EXT2 format output is not directly supported by this program.";
        exit 1
        ;;
    * )
        echo "Unsupported output format: '$OUTPUT_FORMAT'.";
        print_usage
        ;;
esac

if [ -z $OUTPUT_NAME ]; then
    echo "You must specify an output file name with the -o option."
    echo
    print_usage
    exit 1
fi

# ask before overwriting output target
if [ -d $OUTPUT_NAME -o -f $OUTPUT_NAME ]; then
    echo "Output target '$OUTPUT_NAME' will be overwritten by this operation! Are you sure?"
    select yn in "Yes" "No"; do
	case $yn in
	    Yes ) DELETE_OUTPUT=true ; break ;;
	    No  ) echo "Aborting $(basename $0)"; exit 0 ;;
	esac
    done
fi

if [ "$VERBOSE" = "false" ]; then
    DOCKER_BUILD_MODE="--quiet";
fi

# Retry removing Docker network
docker_network_remove()
{
    i=0
    until [[ $i -ge $DOCKER_RETRIES ]]
    do
        # Try to remove docker network
        if (docker network remove $NETWORK_NAME)
        then
            rmerrexit=$?
            break
        else
            rmerrexit=$?
            i=$((i+1))
        fi
        # Handle failed removal
        if [[ $rmerrexit -ne 0 ]] && [[ $i -ge $DOCKER_RETRIES ]]
        then
            # Exit with error
            echo "Failed to remove Docker network $NETWORK_NAME"
            if (docker network ls | grep $NETWORK_NAME)
            then
                # Print out network information for debugging
                docker network inspect $NETWORK_NAME
            fi
            exit $rmerrexit
        else
            # Attempt to resolve and retry removal
            if (docker network ls | grep $NETWORK_NAME >> /dev/null)
            then
                # Force disconnect any containers still attached to network
                for CONTAINER in $(docker network inspect --format '{{range $k, $v := .Containers}}{{print $k}}{{end}}' $NETWORK_NAME)
                do
                    echo "Force removing $CONTAINER from $NETWORK_NAME"
                    docker network disconnect -f $NETWORK_NAME $CONTAINER || true
                done
            else
                echo "Docker network absent: assuming removed."
                break
            fi
            echo "Docker network retry $i: sleeping $i seconds..."
            sleep $i
        fi
    done
}

# Download the docker image from dockerhub, or build it locally from a docker file.
get_image()
{   
    if [ "$USE_DOCKERFILE" = "true" ]; then
        # in either quiet or none quiet mode, IMAGE id will always be exported to intermediate iidfile
        rm -f $TEMP_IMAGE_IIDFILE
        # Create and use custom network for each image build to avoid collisions
        # Note: By default, Docker network range will only allow for 30 networks at a time or a max of
        # 27 concurent image builds. If more is desired, then Docker's default-address-pool will need
        # to be changed to include more ip addresses.
        docker network create $NETWORK_NAME
        $(dirname $0)/myst-retry docker build ${DOCKER_BUILD_MODE} --iidfile ${TEMP_IMAGE_IIDFILE} -f ${DOCKER_FILE} --network ${NETWORK_NAME} ${EXTRA_ARGS} .
        docker_network_remove
        if [ ! -f "$TEMP_IMAGE_IIDFILE" ]; then
            echo "failed to build from Docker Image file $TEMP_IMAGE_IIDFILE";
            exit 1
        fi
        IMAGE_NAME=$(cat $TEMP_IMAGE_IIDFILE)
        rm -f $TEMP_IMAGE_IIDFILE
        DELETE_IMAGE=true
    else
        docker pull $IMAGE_NAME
	    DELETE_IMAGE=false
    fi
}

export_via_overlayfs_mount()
{
    LOWER_DIR=$(docker image inspect --format='{{json .GraphDriver.Data.LowerDir}}' $IMAGE_NAME | tr -d '"')
    UPPER_DIR=$(docker image inspect --format='{{json .GraphDriver.Data.UpperDir}}' $IMAGE_NAME | tr -d '"')
    WORK_DIR=$(docker image inspect --format='{{json .GraphDriver.Data.WorkDir}}' $IMAGE_NAME | tr -d '"')
    MNT_PNT_DIR=$(mktemp -d -t myst-mntpnt.XXXXXX)

    sudo mount -t overlay overlay -olowerdir=$LOWER_DIR,upperdir=$UPPER_DIR,workdir=$WORK_DIR $MNT_PNT_DIR
    # change ownership back to current effective user and group
    sudo chown -R $USER:$(id -ng) $MNT_PNT_DIR

    # copy appenv.json
    mv $TEMP_APPENV $MNT_PNT_DIR/appenv.json

    if [ "$OUTPUT_FORMAT" == "dir" ]; then
        cp -r $MNT_PNT_DIR $OUTPUT_NAME
    elif [ "$OUTPUT_FORMAT" == "cpio" ]; then
        myst mkcpio $MNT_PNT_DIR $OUTPUT_NAME
    fi

    sudo umount $MNT_PNT_DIR
    rm -rf $MNT_PNT_DIR
}

export_via_docker_export()
{
    EXPORTED_IMG_TAR=$(mktemp -t myst-img-tar.XXXXXX)
    EXTRACTED_ROOTFS_DIR=$(mktemp -d -t myst-rootfs-dir.XXXXXX)

    # note that we have to start up the image to flatten it
    TEMP_INAME=$(docker run -d $IMAGE_NAME)
    docker stop $TEMP_INAME >/dev/null
    docker export $TEMP_INAME -o $EXPORTED_IMG_TAR
    docker rm $TEMP_INAME >/dev/null
    if [ -z $DELETE_IMAGE ]; then
        docker rmi -f $IMAGE_NAME
    fi

    # extract rootfs directory from image tarball
    tar xf $EXPORTED_IMG_TAR -C $EXTRACTED_ROOTFS_DIR
    rm -rf $EXPORTED_IMG_TAR

    # copy appenv.json
    mv $TEMP_APPENV $EXTRACTED_ROOTFS_DIR/appenv.json

    if [ "$OUTPUT_FORMAT" == "dir" ]; then
        mv $EXTRACTED_ROOTFS_DIR $OUTPUT_NAME
    elif [ "$OUTPUT_FORMAT" == "cpio" ]; then
        myst mkcpio $EXTRACTED_ROOTFS_DIR $OUTPUT_NAME
    fi

    rm -rf $EXTRACTED_ROOTFS_DIR
}

export_image()
{
    TEMP_APPENV="appenv.json-$(date +%s)"
    docker image inspect --format='{{json .Config}}' $IMAGE_NAME > $TEMP_APPENV

    # cleanup $OUTPUT_NAME if it exists
    rm -rf $OUTPUT_NAME

    if [ "$USE_MOUNT" = "true" ]; then
        export_via_overlayfs_mount
    else
        export_via_docker_export
    fi
}

check_symbols() {
    $(dirname $0)/symbol_check/check_unsupported_symbols.sh $IMAGE_NAME "$PWD/$OUTPUT_NAME"
}

get_image
export_image
if [ "$CHECK_SYMBOL" = "true" ]; then
    check_symbols
fi

echo "Success! Application built at ${OUTPUT_NAME}."
