#!/bin/bash

#
# Copyright (c) 2023-2026, Cloudera, Inc. All Rights Reserved.
#
# 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
#
#     http://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.
#

set -euo pipefail

# Default values
TRUSTSTORE_FORMAT="PKCS12"
TRUSTSTORE_PASSWORD=""
ALIAS="server-cert"
AUTO_YES=false
SOURCE_TRUSTSTORE=""
SOURCE_TRUSTSTORE_PASSWORD=""
SOURCE_TRUSTSTORE_TYPE="JKS"

# Usage function
usage() {
    cat << EOF
Usage: $0 -s|--ssl-server <host:port> -t|--truststore <truststore-path> [-p|--password <password>] [-f|--format <format>] [-a|--alias <alias>] [-y|--yes]

Downloads SSL certificate from an SSL server using Java keytool and stores it in a truststore.
If the truststore does not exist, it will be created.

Options:
    -s, --ssl-server    SSL server in host:port format (required, e.g., cm-host:7180)
    -t, --truststore    Path to truststore file (required)
    -p, --password      Truststore password (required)
    -f, --format        Truststore format: PKCS12 or JKS (default: PKCS12)
    -a, --alias         Certificate alias in truststore (default: server-cert)
    -y, --yes           Automatically overwrite existing aliases without prompting
    -j, --java-home     Path to Java home directory (default: \$JAVA_HOME)
    -h, --help          Show this help message

Examples:
    $0 -s cm-host:7180 -t /etc/cma/truststore.p12
    $0 --ssl-server cm-host:7443 --truststore /tmp/truststore.p12 --password mypassword
    $0 -s server:8443 -t truststore.jks -f JKS -a my-server --yes

EOF
    exit 1
}

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        -s|--ssl-server)
            SSL_SERVER="$2"
            shift 2
            ;;
        -t|--truststore)
            TRUSTSTORE_PATH="$2"
            shift 2
            ;;
        -p|--password)
            TRUSTSTORE_PASSWORD="$2"
            shift 2
            ;;
        -f|--format)
            TRUSTSTORE_FORMAT="$2"
            shift 2
            ;;
        -a|--alias)
            ALIAS="$2"
            shift 2
            ;;
        -y|--yes)
            AUTO_YES=true
            shift
            ;;
        --source-truststore)
            SOURCE_TRUSTSTORE="$2"
            shift 2
            ;;
        --source-truststore-password)
            SOURCE_TRUSTSTORE_PASSWORD="$2"
            shift 2
            ;;
        --source-truststore-type)
            SOURCE_TRUSTSTORE_TYPE="$2"
            shift 2
            ;;
        -j|--java-home)
            JAVA_HOME="$2"
            shift 2
            ;;
        -h|--help)
            usage
            ;;
        *)
            echo "Unknown option: $1"
            usage
            ;;
    esac
done

KEYTOOL="${JAVA_HOME:+${JAVA_HOME}/bin/}keytool"

# Validate required parameters
if [ -z "$SSL_SERVER" ]; then
    echo "Error: SSL server is required"
    usage
fi

if [ -z "$TRUSTSTORE_PATH" ]; then
    echo "Error: Truststore path is required"
    usage
fi

if [ -z "$TRUSTSTORE_PASSWORD" ]; then
    echo "Error: Truststore password is required"
    usage
fi

# Validate truststore format
if [ "$TRUSTSTORE_FORMAT" != "PKCS12" ] && [ "$TRUSTSTORE_FORMAT" != "JKS" ]; then
    echo "Error: Truststore format must be PKCS12 or JKS"
    exit 1
fi

# Parse host and port from SSL_SERVER
if [[ "$SSL_SERVER" =~ ^([^:]+):([0-9]+)$ ]]; then
    HOST="${BASH_REMATCH[1]}"
    PORT="${BASH_REMATCH[2]}"
elif [[ "$SSL_SERVER" =~ ^[^:]+$ ]]; then
    HOST="$SSL_SERVER"
    PORT=443
else
    echo "Error: Invalid SSL server format. Use host:port or just host (default port 443)"
    exit 1
fi

echo "Downloading SSL certificate from: $HOST:$PORT"
echo "Truststore: $TRUSTSTORE_PATH"
echo "Format: $TRUSTSTORE_FORMAT"
echo "Alias: $ALIAS"
echo ""

# Create directory for truststore if it doesn't exist
TRUSTSTORE_DIR=$(dirname "$TRUSTSTORE_PATH")
if [ ! -d "$TRUSTSTORE_DIR" ]; then
    echo "Creating directory: $TRUSTSTORE_DIR"
    mkdir -p "$TRUSTSTORE_DIR"
fi

# Download the full certificate chain using keytool
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
TEMP_CHAIN="/tmp/cert_chain_${TIMESTAMP}_$$.pem"
trap "rm -f '${TEMP_CHAIN}' /tmp/cert_${TIMESTAMP}_$$_*.pem" EXIT

echo "Downloading certificate chain from $HOST:$PORT..."
KEYTOOL_OPTS=()
if [ -n "${SOURCE_TRUSTSTORE}" ]; then
    echo "Using source truststore: ${SOURCE_TRUSTSTORE} (type: ${SOURCE_TRUSTSTORE_TYPE})"
    KEYTOOL_OPTS+=(
        "-J-Djavax.net.ssl.trustStore=${SOURCE_TRUSTSTORE}"
        "-J-Djavax.net.ssl.trustStoreType=${SOURCE_TRUSTSTORE_TYPE}"
    )
    if [ -n "${SOURCE_TRUSTSTORE_PASSWORD}" ]; then
        KEYTOOL_OPTS+=("-J-Djavax.net.ssl.trustStorePassword=${SOURCE_TRUSTSTORE_PASSWORD}")
    fi
fi
KEYTOOL_ERR=$(mktemp)
if ! timeout 10 "${KEYTOOL}" ${KEYTOOL_OPTS[@]+"${KEYTOOL_OPTS[@]}"} -printcert -sslserver "$HOST:$PORT" -rfc > "$TEMP_CHAIN" 2>"$KEYTOOL_ERR" || [ ! -s "$TEMP_CHAIN" ]; then
    echo "Error: Failed to download certificate from $HOST:$PORT"
    echo "Please verify the host, port, and SSL/TLS availability"
    [ -s "$KEYTOOL_ERR" ] && echo "keytool output:" && cat "$KEYTOOL_ERR"
    rm -f "$KEYTOOL_ERR"
    exit 1
fi
rm -f "$KEYTOOL_ERR"

# Split the chain PEM into individual certificate files
awk '/-----BEGIN CERTIFICATE-----/{idx++; file=sprintf("/tmp/cert_'${TIMESTAMP}'_'$$'_%02d.pem", idx-1)}
     idx{print > file}
     /-----END CERTIFICATE-----/{if(idx) close(file)}' "$TEMP_CHAIN"

# Collect the split files in order
CERT_FILES=()
for f in $(ls /tmp/cert_${TIMESTAMP}_$$_*.pem 2>/dev/null | sort); do
    CERT_FILES+=("$f")
done

CERT_COUNT=${#CERT_FILES[@]}
if [ "$CERT_COUNT" -eq 0 ]; then
    echo "Error: No certificates found in response from $HOST:$PORT"
    exit 1
fi

echo "Certificate chain depth: $CERT_COUNT"

# Build the list of aliases to import (single cert uses base alias; chain uses indexed aliases)
declare -a ALIASES
for (( i=0; i<CERT_COUNT; i++ )); do
    if [ "$CERT_COUNT" -eq 1 ]; then
        ALIASES+=("$ALIAS")
    else
        ALIASES+=("${ALIAS}-${i}")
    fi
done

# Check if any target alias already exists in the truststore and prompt once
EXISTING_ALIASES=()
if [ -f "$TRUSTSTORE_PATH" ]; then
    for a in "${ALIASES[@]}"; do
        if "${KEYTOOL}" -list -keystore "$TRUSTSTORE_PATH" -storepass "$TRUSTSTORE_PASSWORD" -storetype "$TRUSTSTORE_FORMAT" -alias "$a" &>/dev/null; then
            EXISTING_ALIASES+=("$a")
        fi
    done
fi

if [ ${#EXISTING_ALIASES[@]} -gt 0 ]; then
    echo "Warning: The following aliases already exist in the truststore: ${EXISTING_ALIASES[*]}"
    if [ "$AUTO_YES" = true ]; then
        echo "Overwriting existing aliases (-y/--yes specified)"
    else
        read -p "Delete and re-import? (y/n) " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            echo "Import cancelled"
            exit 0
        fi
    fi
    for a in "${EXISTING_ALIASES[@]}"; do
        "${KEYTOOL}" -delete -keystore "$TRUSTSTORE_PATH" -storepass "$TRUSTSTORE_PASSWORD" -storetype "$TRUSTSTORE_FORMAT" -alias "$a"
    done
    echo "Existing certificates deleted"
fi

# Import each certificate in the chain
echo "Importing certificate chain into truststore..."
for (( i=0; i<CERT_COUNT; i++ )); do
    CERT_ALIAS="${ALIASES[$i]}"
    LABEL=$([ "$i" -eq 0 ] && echo "leaf" || echo "intermediate/CA $i")
    echo "  Importing cert $((i+1))/$CERT_COUNT ($LABEL) as alias '$CERT_ALIAS'..."
    "${KEYTOOL}" -importcert -noprompt \
        -keystore "$TRUSTSTORE_PATH" \
        -storepass "$TRUSTSTORE_PASSWORD" \
        -storetype "$TRUSTSTORE_FORMAT" \
        -alias "$CERT_ALIAS" \
        -file "${CERT_FILES[$i]}"
done

echo "Certificate chain imported successfully"
echo ""
echo "Truststore entries:"
for a in "${ALIASES[@]}"; do
    "${KEYTOOL}" -list -keystore "$TRUSTSTORE_PATH" -storepass "$TRUSTSTORE_PASSWORD" -storetype "$TRUSTSTORE_FORMAT" -alias "$a" 2>/dev/null
done
