#
# Generate compose yml for specific level.
#
# Arguments:
#   - $1 = level: Create lab for specific level.
#   - $2 = red_team_services: Spawn additional red team services.
#   - $3 = blue_team_services: Spawn additional blue team services.
#   - $4 = monitoring_services: Spawn additional monitring services.
#   - $5 = all_services: If set to 1 spawn all available services.
#
# Returns:
#   - 0: Always.
#
# Prints:
#   - stdout: Nothing.
#   - stderr: Error message if a given service class has not config files.
#
# Creates:
#   - Files:
#     - ./docker-compose.yml: Config file for docker-compose to start lab.
#
_labctl_buildComposeFile() {
    local level_array="$1"
    local red_team_services="$2"
    local blue_team_services="$3"
    local monitoring_services="$4"
    local all_services="$5"

    local class
    local services
    local config_path
    local service
    # Variable is defined in ../lab.sh.
    # shellcheck disable=SC2154
    local victim_path="${working_dir}/etc/services/victim"
    local level
    local volumes_tmp_file
    local volumes

    {
    printf '# Do not edit this file manually.\n'
    printf '# This file is auto generated by ./lab.sh.\n'
    printf '# To change this file edit files located in: "./etc/services/".\n'
    printf '# Afterwards restart the lab.\n'
    } > "${working_dir}/docker-compose.yml"

    cat "${working_dir}/etc/services/default.yml"\
        >> "${working_dir}/docker-compose.yml"

    sed -i "s/nope_uid/$(id -u)/g" "${working_dir}/docker-compose.yml"
    sed -i "s/nope_gid/$(id -g)/g" "${working_dir}/docker-compose.yml"

    if [ "${all_services}" -eq 1 ]; then
        find "${working_dir}/etc/services/"\
            -mindepth 2 -type f -name '*.yml'\
            -exec sed -e 's/^/  /' {} \;\
            >> "${working_dir}/docker-compose.yml"
    else
        for class in red_team blue_team monitoring; do
            case "${class}" in
                red_team)
                    services="${red_team_services}"
                    config_path="${working_dir}/etc/services/red_team"
                    ;;
                blue_team)
                    services="${blue_team_services}"
                    config_path="${working_dir}/etc/services/blue_team"
                    ;;
                monitoring)
                    services="${monitoring_services}"
                    config_path="${working_dir}/etc/services/monitoring"
                    ;;
            esac

            for service in ${services}; do
                if  [ "${service}" = 'all' ]; then
                    if [ "$(find "${config_path}"\
                        -type f -name '*.yml' | wc -l)" -eq 0 ]; then
                        printf 'Error: No %s configs found.\n' "${class}" >&2
                        exit 1
                    else
                        find "${config_path}"\
                            -type f -name '*.yml'\
                            -exec sed -e 's/^/  /' {} \;\
                            >> "${working_dir}/docker-compose.yml"
                    fi
                elif ! [ "${service}" = 'none' ]; then
                    if [ "$(find "${config_path}"\
                        -type f -name "${service}.yml" | wc -l)" -ne 1 ]; then
                        printf 'Error: No config found for %s service: %s.\n'\
                            "${class}" "${service}" >&2
                        exit 1
                    else
                        find "${config_path}"\
                            -type f -name "${service}.yml"\
                            -exec sed -e 's/^/  /' {} \;\
                            >> "${working_dir}/docker-compose.yml"
                    fi
                fi
            done
        done

        for level in ${level_array}; do
            case "${level}" in
                all)
                    find "${working_dir}/etc/services/victim/"\
                        -type f -name '*.yml'\
                        -exec sed -e 's/^/  /' {} \;\
                            >> "${working_dir}/docker-compose.yml"
                    ;;
                *)
                    if [ "$(find "${victim_path}/${level}" -type f -name '*.yml'\
                        | wc -l)" -eq 0 ]; then
                        printf 'No services found for level: %s.\n' "${level}"
                        exit 1
                    else
                        find "${victim_path}/${level}" -type f -name '*.yml'\
                            -exec sed -e 's/^/  /' {} \;\
                            >> "${working_dir}/docker-compose.yml"
                    fi
                    ;;
            esac
        done
    fi

    printf '\n' >> "${working_dir}/docker-compose.yml"

    # Generate volume list.
    volumes_tmp_file="$(mktemp)"
    volumes="$(yq '.services | .[] | .volumes'\
        "${working_dir}/docker-compose.yml" -crM\
        | tr -d '\["\]' | cut -d ':' -f 1 | sed '/^null$/d;/^.\//d;/^\//d')"

    printf '%s\n' 'volumes:' > "${volumes_tmp_file}"

    while read -r volume; do 
        printf '  %s:\n' "${volume}" >> "${volumes_tmp_file}"
    done <<<"${volumes}"

    # Add volume list to docker-compose.yml.
    cat "${volumes_tmp_file}" >> "${working_dir}/docker-compose.yml"
    rm "${volumes_tmp_file}"
}

#
# Check for existing ssh keys.
# If no keys present exit.
# If keys found add these to kali box, for ssh public key authentication.
#
# Arguments:
#   - None
#
# Returns:
#   - 0: Always.
#
# Prints:
#   - stdout: Nothing.
#   - stderr: Error message if no SSH keys found.
#
# Creates:
#   - Files:
#     - ./etc/keys: Directory containing SSH public keys.
#     - ./etc/kali/authorized_keys: File containing public keys to add to kali box.
#
_labctl_getKeys() {
    if [ "$(find "$HOME/.ssh/" -name '*.pub' | wc -l)" -eq 0 ]; then
        printf 'Error: No SSH public keys found in %s.\n\n' "$HOME/.ssh/" >&2
        printf 'Please generate SSH keys. For example: %s.\n'\
            'ssh-keygen -ted25519' >&2
        printf 'Otherwise you won'\''t be able to use the lab.\n' >&2
        exit 1
    fi
 
    _labctl_removeKeys

    # Variable is defined in ../lab.sh.
    # shellcheck disable=SC2154
    if ! [ -d "${key_dir}" ]; then
        mkdir -p "${key_dir}"
    fi
    
    find "$HOME/.ssh/" -name '*.pub' -exec cp {} "${key_dir}" \;
    # Variable is defined in ../lab.sh.
    # shellcheck disable=SC2154
    find "${key_dir}" -name '*.pub' -exec cat {} \; >> "${auth_keys}"
}

#
# Prepare needed ressources for lab to work.
#
# Arguments:
#   - None
#
# Returns:
#   - 0: Always.
#
# Prints:
#   - stdout: Nothing.
#   - stderr: Nothing.
#
# Creates:
#   - Nothing
#
_labctl_prepare() {
    # Install loki log driver for docker.
    if ! docker plugin list | grep -q 'loki:latest'; then
        docker plugin \
            install \
            grafana/loki-docker-driver:latest \
            --alias loki \
            --grant-all-permissions || exit 1
    fi

    tar xf ./etc/heimdall.tar ./etc/heimdall
}

#
# Build all needed ressources for the lab.
#
# Arguments:
#   - None
#
# Returns:
#   - 0: Always.
#
# Prints:
#   - stdout: Nothing.
#   - stderr: Nothing.
#
# Creates:
#   - Nothing
#
_labctl_build() {
    _labctl_getKeys
    _labctl_prepare
    docker-compose build --parallel || exit 1 
}

#
# Call function to generate compose yml and start the lab afterwards.
#
# Arguments:
#   - $1 = level: Create lab for specific level.
#   - $2 = red_team_services: Spawn additional red team services.
#   - $3 = blue_team_services: Spawn additional blue team services.
#   - $4 = monitoring_services: Spawn additional monitring services.
#   - $5 = all_services: If set to 1 spawn all available services.
#
# Returns:
#   - 0: Always.
#
# Prints:
#   - stdout: Nothing.
#   - stderr: Nothing.
#
# Creates:
#   - Files:
#     - $HOME/.local/var/run/pentest_lab/info: Persistently save lab state.
#
_labctl_up() {
    local level="$1"
    local red_team_services="$2"
    local blue_team_services="$3"
    local monitoring_services="$4"
    local all_services="$5"
    local run_dir

    run_dir="$HOME/.local/var/run/pentest_lab"

    mkdir -p "${run_dir}" || true
    printf '' > "${run_dir}/info"

    {
    printf 'level: %s\n' "${level}"
    printf 'red_team_services: %s\n' "${red_team_services}"
    printf 'blue_team_services: %s\n' "${blue_team_services}"
    printf 'monitoring_services: %s\n' "${monitoring_services}"
    } >> "${run_dir}/info"

    _labctl_buildComposeFile\
        "${level}"\
        "${red_team_services}"\
        "${blue_team_services}"\
        "${monitoring_services}"\
        "${all_services}"
    _labctl_build
    docker-compose up -d --remove-orphans || _labctl_emergencyCleanup
}

#
# Remove orphaned files if lab startup fails.
#
# Arguments:
#   - None
#
# Returns:
#   - 0: Always.
#
# Prints:
#   - stdout: Nothing.
#   - stderr: Nothing.
#
# Creates:
#   - Nothing
#
_labctl_emergencyCleanup() {
    local run_dir
    run_dir="$HOME/.local/var/run/pentest_lab"


    if [ -d "${run_dir}" ]; then
        rm -rf "${run_dir}"
    fi

    docker-compose down --remove-orphans

    printf '\n\e[31m\e[1mSome error occured.\nEmergency cleanup done.\n\e[0m' \
        >&2
    exit 1
}

#
# Stop the lab.
#
# Arguments:
#   - None
#
# Returns:
#   - 0: Always.
#
# Prints:
#   - stdout: Nothing.
#   - stderr: Nothing.
#
# Creates:
#   - Nothing
#
_labctl_down() {
    local run_dir

    run_dir="$HOME/.local/var/run/pentest_lab"

    docker-compose down --remove-orphans

    if [ -d "${run_dir}" ]; then
        rm -rf "${run_dir}"
    fi

    tar cf ./etc/heimdall.tar ./etc/heimdall
}

#
# Remove all ressources owned by the lab - reset the lab.
#
# Arguments:
#   - None
#
# Returns:
#   - 0: Always.
#
# Prints:
#   - stdout: Nothing.
#   - stderr: Nothing.
#
# Creates:
#   - Nothing
#
_labctl_prune() {
    docker-compose down --remove-orphans --rmi all --volumes
    docker-compose rm -fsv
    _labctl_removeKeys
}

#
# Remove ssh keys from local directory.
#
# Arguments:
#   - None
#
# Returns:
#   - 0: Always.
#
# Prints:
#   - stdout: Nothing.
#   - stderr: Nothing.
#
# Creates:
#   - Nothing
#
_labctl_removeKeys() {
    printf 'Removing: %s.\n' "${key_dir}"
    if [ -d "${key_dir}" ]; then
        rm -r "${key_dir}"
    fi

    printf 'Removing: %s.\n' "${auth_keys}"
    if [ -f "${auth_keys}" ]; then
        rm "${auth_keys}"
    fi
}
