.kommands

## kommands - k8s commands
## https://github.com/javimox/kommands
## v0.1 - 2020.06.21
## v0.2 - 2021.02.22 - add ingress suport
## v0.3 - 2021.02.24 - add cm and secret support
## v2   - 2023.01.24 - read namespace from parameter
##
##  ============================================================================
##  Copyright (C) 2023 jmox@pm.me
##
##  This program is free software: you can redistribute it and/or modify it
##  under the terms of the GNU General Public License as published by the Free
##  Software Foundation, either version 3 of the License, or (at your option)
##  any later version.
##
##  This program is distributed in the hope that it will be useful, but WITHOUT
##  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
##  FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
##  more details.
##
##  You should have received a copy of the GNU General Public License along with
##  this program.  If not, see <http://www.gnu.org/licenses/>.
##  ============================================================================
##

## ------------------
## helpers
## ------------------

_shift() {
    echo "$@" | sed -e "s/${1}//g" | xargs
}

_find_namespace() {
    sed '/\(^\| \)-n */!d; s//\n/; s/^.*\n//; s/ .*//' <<< "$@"
}

_sort_args() {
    if [ -z "$_k_ns" ] ; then
        [[ -z "$_k_args" ]] && declare -g _k_args=""
        declare -g _k_ns=""
        _k_ns=$(_find_namespace $@)
        if [ ! -z "$_k_ns" ] ; then
            [[ "$@" =~ "-n${_k_ns}" ]] && _k_ns="-n${_k_ns}" || _k_ns="-n ${_k_ns}"
            _k_args=$(_shift "${_k_ns}" "$@" | xargs)
        else
            _k_args="$@"
        fi
    fi
}

# check if string contains only numbers
_is_nr() {
    [[ $1 =~ ^[0-9]+$ ]] && return 0 || return 1
}

# add prefix line number to output
_add_line_nr() {
    awk '{printf "%2d %s\n", NR-1, $0}'
}

# add prefix number to autocomplete words
_add_autocomplete_nr() {
    local c=1
    local total=$#
    local option=""
    for (( c=1; c<=$total; c++ )) ; do
        printf "$c:${!c} "
    done
}

# get name from line number
_convert_k8s_nr2name() {
    local line_requested=$(( $1 + 1 ))
    awk "NR==$line_requested" | awk '{print $2}'
}

## show all pod,pvc,svc,ingress or one specificied by name or number
_k_get() {
    # $1 resource type, $2=_k_args resource number or name , $3=_k_ns namespace
    local run=""
    local rt="$1"       # resource type
    local rn="$2"       # resource name/nr
    local ns="$3"       # namespace

    # empty global vars
    _k_args=""
    _k_ns=""

    if [ -z "$rn" ] ; then
        # resource name/nr not specified, listing all
        run='kubectl get "$rt" ${ns//[[:space:]]:+"$ns"} | _add_line_nr'
    else
        if _is_nr "$rn" ; then
            local line_requested=$(( $rn + 1 ))
            run='kubectl get "$rt" ${ns//[[:space:]]:+"$ns"} | _add_line_nr | sed -n "1p;${line_requested}p"'
        else
            run='kubectl get "$rt" ${rn//[[:space:]]:+"$rn"} ${ns//[[:space:]]:+"$ns"} | _add_line_nr'
        fi
    fi

    eval "${run}"
}

# delete pod,pvc,svc,ingress using their name or their number
_k_del() {
    # $1 resource type, $2=_k_args resource number or name , $3=_k_ns namespace
    local rt="$1"       # resource type
    local rn="$2"       # resource name/nr
    local ns="$3"       # namespace

    # empty global vars
    _k_args=""
    _k_ns=""

    if [ -z "$rn" ] ; then
        # resource name/nr not specified, param mandatory
        echo "error delete: required object number or name not specified."
    else
        if _is_nr "$rn" ; then
            local run=""
            local line_requested=$(( $rn + 1 ))
            case "$rt" in
                po*) run="kgpo";; pvc) run="kgpvc";; svc) run="kgsvc";;
                 cm) run="kgcm";; in*) run="kgin";; se*) run="kgse";; *) ;;
            esac
            run="${run} ${ns//[[:space:]]:+\"$ns\"}"
            rn=$(eval "${run}" | _convert_k8s_nr2name $rn)
        fi

        if [ -z "$rn" ] ; then
            echo "error delete: $rt not found"
        else
            read -p "do you want to remove ${rn} (y/N)? " choice
            case "$choice" in
                y|Y ) kubectl delete "${rt}" "${rn}" ${ns//[[:space:]]:+"$ns"} ;;
                *) ;;
            esac         
        fi
    fi
}

# show logs -f of a pod / container
_k_log() {
    # $2=_k_args resource number or name , $3=_k_ns namespace
    local ns="$3"       # namespace
    local rn_arr=()
    local pod=""
    local cont=""

    read -a rn_arr <<< "$2"
    pod=${rn_arr[0]}
    cont=${rn_arr[1]}

    # empty global vars
    _k_args=""
    _k_ns=""

    if [ -z "$pod" ] ; then
        # resource name/nr not specified, param mandatory
        echo "error show log: required object number or name not specified."
    else
        local podname=""
        local container=""
        _is_nr "$pod" && podname="$(kgpo ${ns//[[:space:]]:+\"$ns\"} | _convert_k8s_nr2name $pod)" || podname="$pod"
        if [ ! -z "$cont" ] ; then
            _is_nr "$cont" && container="$(ksco ${podname} ${ns//[[:space:]]:+\"$ns\"} | _convert_k8s_nr2name $cont)" || container="$cont"
            container=( -c ${container} )
        fi
        [[ -z "$ns" ]] && ns="-n $(kubectl config view --minify -o jsonpath='{..namespace}')"
        echo "(${ns:=-n default}) Showing logs of: ${podname} ${container[@]}"
        kubectl logs -f ${podname} "${container[@]}" ${ns//[[:space:]]:+"$ns"}
    fi
}

# show pod's containers
_k_sco() {
    # $2=_k_args resource number or name , $3=_k_ns namespace
    local rn="$2"       # resource name/nr
    local ns="$3"       # namespace

    # empty global vars
    _k_args=""
    _k_ns=""

    if [ "$#" -gt "0" ] ; then
        local podname=""
        _is_nr "$rn" && podname="$(kgpo ${ns//[[:space:]]:+\"$ns\"} | _convert_k8s_nr2name $rn)" || podname="$rn"
        local line="$(printf "pod:${podname} " ; kubectl get pods "${podname}" -o jsonpath='{.spec.containers[*].name}' ${ns//[[:space:]]:+\"$ns\"})"
        echo "$line" | tr ' ' '\n' | _add_line_nr
    else
        echo "error show containers: required pod number or name not specified."
    fi
}

# describe resource
_k_desc() {
    local run=""
    local ns="$2"       # namespace
    local r_type=""     # provided by user
    local r_name=""     # provided by user
    local type=""       # resource type
    local resource_name="" # resource name

    read -a rn_arr <<< "$1"
    r_type=${rn_arr[0]}
    r_name=${rn_arr[1]}

    # empty global vars
    _k_args=""
    _k_ns=""

    case "$r_type" in
        po*) run="kgpo" ; type="pods";;
        pvc) run="kgpvc"; type="pvc" ;;
        svc) run="kgsvc"; type="svc" ;;
        cm)  run="kgcm" ; type="cm"  ;;
        in*) run="kgin" ; type="ingress" ;;
        se*) run="kgse" ; type="secret"  ;;
        *)   run="kgpo" ; type="pods";;
    esac

    run="${run} ${ns//[[:space:]]:+\"$ns\"}"

    _is_nr "$r_name" && resource_name=$(eval "${run}" | _convert_k8s_nr2name $r_name) || resource_name="$r_name"

    echo "(${ns:=-n default})"
    if [ -z "${resource_name}" ] ; then
        kubectl describe "${type}" ${ns//[[:space:]]:+"$ns"}
    else
        kubectl describe "${type}" "${resource_name}" ${ns//[[:space:]]:+"$ns"}
    fi
}

# open a shell
_k_esh() {
    local ns="$2"       # namespace

    read -a rn_arr <<< "$1"
    pod=${rn_arr[0]}
    cont=${rn_arr[1]}

    # empty global vars
    _k_args=""
    _k_ns=""

    if [ -z "${rn_arr[0]}" ] ; then
        # resource name/nr not specified, param mandatory
        echo "error on shell: required object number or name not specified."
    else
        local podname=""
        local container=""
        _is_nr "$pod" && podname="$(kgpo ${ns//[[:space:]]:+\"$ns\"} | _convert_k8s_nr2name $pod)" || podname="$pod"
        if [ ! -z "$cont" ] ; then
            _is_nr "$cont" && container="$(ksco ${podname} ${ns//[[:space:]]:+\"$ns\"} | _convert_k8s_nr2name $cont)" || container="$cont"
            container=( -c ${container} )
        fi
        [[ -z "$ns" ]] && ns="-n $(kubectl config view --minify -o jsonpath='{..namespace}')"
        echo "(${ns:=-n default}) You are now in pod: ${podname} ${container[@]}"
        kubectl exec -it "${podname}" "${container[@]}" ${ns//[[:space:]]:+"$ns"} -- sh
    fi
}

# autocomplete handler
_k_autocomplete() {
    local cur=${COMP_WORDS[COMP_CWORD]}
    local cmd="${1##*/}"
    local action="get"    # default
    local resource="pod"  # default
    local opts=""         # options output
    local parse=0         # add number to output

    case "${cmd}" in
        kgpo|kdpo|ksco) resource="pod" ;;
        kgpvc|kdpvc)    resource="pvc" ;;
        kgsvc|kdsvc)    resource="svc" ;;
        kgcm|kdcm)       resource="cm" ;;
        kgin|kdin)  resource="ingress" ;;
        kgse|kdse)   resource="secret" ;;
        kesh|klog)       action="exec" ;;
        kdesc)       action="describe" ;;
        *) ;;
    esac

    case "${action}" in
        get)
            opts="$(kubectl ${action} ${resource} -o=jsonpath='{range .items[*].metadata.name}{@}{"\n"}{end}')"
        ;;
        exec)
            case ${COMP_CWORD} in
                1) opts="$(kubectl get pod -o=jsonpath='{range .items[*].metadata.name}{@}{"\n"}{end}')" ;;
                2) opts="$(kubectl get pod -o=jsonpath='{.spec.containers[*].name}' ${3})" ;;
                *) ;;
            esac
        ;;
        describe)
            case ${COMP_CWORD} in
                1) opts="$(kubectl api-resources -o name --cached --request-timeout=5s --verbs=get) pvc svc" ; parse=1 ;;
                2) opts="$(kubectl get -o template '--template={{ range .items  }}{{ .metadata.name }} {{ end }}' ${3})" ;;
                *) ;;
            esac
        ;;
        *) ;;
    esac

    # add number to autocompletion output. If enabled, it won't autocomplete using names
    # [[ "$parse" ]] && opts="$(_add_autocomplete_nr ${opts})"

    COMPREPLY=( $(compgen -W  "${opts}" -- ${cur}) )
    return 0
}

# autocomplete yaml|yml|json files and outputs directories
_kapp_autocomplete ()
{
    local IFS=$'\n'
    local last_char=' '

    COMPREPLY=($(compgen -o plusdirs -f -X '!(*.yaml|*.yml|*.json)' \
        -- "${COMP_WORDS[COMP_CWORD]}"))

    if [ ${#COMPREPLY[@]} = 1 ] ; then
        [ -d "${COMPREPLY}" ] && last_char=/
        COMPREPLY=$(printf %q%s "${COMPREPLY}" "${last_char}")
    else
        for ((i=0; i < ${#COMPREPLY[@]}; i++)); do
            [ -d "${COMPREPLY[$i]}" ] && COMPREPLY[$i]=${COMPREPLY[$i]}/
        done
    fi
    return 0
}

## ------------------
## k commands
## ------------------

# get pods
kgpo() {
    _sort_args "$@"
    _k_get "pods" "$_k_args" "$_k_ns"
}

# get persistentVolumeClaims
kgpvc() {
    _sort_args "$@"
    _k_get "pvc" "$_k_args" "$_k_ns"
}

# get services
kgsvc() {
    _sort_args "$@"
    _k_get "svc" "$_k_args" "$_k_ns"
}

# get configmaps
kgcm() {
    _sort_args "$@"
    _k_get "cm" "$_k_args" "$_k_ns"
}

# get ingress
kgin() {
    _sort_args "$@"
    _k_get "ingress" "$_k_args" "$_k_ns"
}

# get secrets
kgse() {
    _sort_args "$@"
    _k_get "secret" "$_k_args" "$_k_ns"
}

# delete pod
kdpo() {
    _sort_args "$@"
    _k_del "pod" "$_k_args" "$_k_ns"
}

# delete pvc
kdpvc() {
    _sort_args "$@"
    _k_del "pvc" "$_k_args" "$_k_ns"
}

# delete service
kdsvc() {
    _sort_args "$@"
    _k_del "svc" "$_k_args" "$_k_ns"
}

# delete configmap
kdcm() {
    _sort_args "$@"
    _k_del "cm" "$_k_args" "$_k_ns"
}

# delete ingress
kdin() {
    _sort_args "$@"
    _k_del "ingress" "$_k_args" "$_k_ns"
}

# delete secret
kdse() {
    _sort_args "$@"
    _k_del "secret" "$_k_args" "$_k_ns"
}

# apply manifest
kapp() {
    kubectl apply -f "$@"
}

# show logs
klog() {
    _sort_args "$@"
    _k_log "pod" "$_k_args" "$_k_ns"
}

# show containers of a pod
ksco() {
    _sort_args "$@"
    _k_sco "pod" "$_k_args" "$_k_ns"
}

# describe resource
kdesc() {
    _sort_args "$@"
    _k_desc "$_k_args" "$_k_ns"
}

# open a shell in the pod/container
kesh() {
    _sort_args "$@"
    _k_esh "$_k_args" "$_k_ns"
}


# enable autocomplete
complete -F _k_autocomplete kgpo kgpvc kgsvc kgcm kgin kgse kdpo kdpvc kdsvc kdcm kdin kdse klog kdesc ksco kesh

# enable autocomplete for kapp
complete -o nospace -F _kapp_autocomplete kapp