ssm

#!/usr/bin/env bash

# AWS SSM Session Manager Helper Script
# Repository: https://github.com/ZSoftly/ztiaws

set -e  # Exit on error
set -u  # Exit on undefined variables

# Get SCRIPT_DIR and source utilities
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Source utilities (like colors)
if [ -f "${SCRIPT_DIR}/src/00_utils.sh" ]; then
    # shellcheck source=./src/00_utils.sh
    source "${SCRIPT_DIR}/src/00_utils.sh"
elif [ -f "/usr/local/bin/src/00_utils.sh" ]; then # For system-wide installation
    # shellcheck source=/dev/null
    source "/usr/local/bin/src/00_utils.sh"
else
    # 00_utils.sh is mandatory
    echo "[ERROR] src/00_utils.sh not found. This script requires 00_utils.sh for color definitions and common functions." >&2
    echo "Please ensure src/00_utils.sh is present in the script directory or in /usr/local/bin/src/." >&2
    exit 1
fi

# Initialize logging (console-only by default for SSM, can be enabled with ENABLE_SSM_LOGGING=true env var)
init_logging "ssm" "${ENABLE_SSM_LOGGING:-false}"

# Import regions from 01_regions.sh
if [ -f "${SCRIPT_DIR}/src/01_regions.sh" ]; then
    # shellcheck source=./src/01_regions.sh
    source "${SCRIPT_DIR}/src/01_regions.sh"
elif [ -f "/usr/local/bin/src/01_regions.sh" ]; then
    # shellcheck source=/dev/null
    source "/usr/local/bin/src/01_regions.sh"
else
    echo "Error: Could not find src/01_regions.sh in either ${SCRIPT_DIR}/src/ or /usr/local/bin/src/" >&2
    exit 1
fi



# Import run_command functions
if [ -f "${SCRIPT_DIR}/src/03_ssm_command_runner.sh" ]; then
    # shellcheck source=./src/03_ssm_command_runner.sh
    source "${SCRIPT_DIR}/src/03_ssm_command_runner.sh"
elif [ -f "/usr/local/bin/src/03_ssm_command_runner.sh" ]; then
    # shellcheck source=/dev/null
    source "/usr/local/bin/src/03_ssm_command_runner.sh"
else
    echo "Error: Could not find src/03_ssm_command_runner.sh in either ${SCRIPT_DIR}/src/ or /usr/local/bin/src/" >&2
    exit 1
fi
# Import instance resolver functions
if [ -f "${SCRIPT_DIR}/src/02_ssm_instance_resolver.sh" ]; then
    source "${SCRIPT_DIR}/src/02_ssm_instance_resolver.sh"
fi

# Import file transfer functions
FILE_TRANSFER_AVAILABLE=false
if [ -f "${SCRIPT_DIR}/src/04_ssm_file_transfer.sh" ]; then
    # shellcheck source=./src/04_ssm_file_transfer.sh
    source "${SCRIPT_DIR}/src/04_ssm_file_transfer.sh"
    FILE_TRANSFER_AVAILABLE=true
elif [ -f "/usr/local/bin/src/04_ssm_file_transfer.sh" ]; then
    # shellcheck source=/dev/null
    source "/usr/local/bin/src/04_ssm_file_transfer.sh"
    FILE_TRANSFER_AVAILABLE=true
else
    log_warn "File transfer module not found. Upload/download commands will not be available."
fi

# Import parameter parser module
if [ -f "${SCRIPT_DIR}/src/07_ssm_parameter_parser.sh" ]; then
    # shellcheck source=./src/07_ssm_parameter_parser.sh
    source "${SCRIPT_DIR}/src/07_ssm_parameter_parser.sh"
elif [ -f "/usr/local/bin/src/07_ssm_parameter_parser.sh" ]; then
    # shellcheck source=/dev/null
    source "/usr/local/bin/src/07_ssm_parameter_parser.sh"
else
    echo "[ERROR] src/07_ssm_parameter_parser.sh not found. This script requires the parameter parser module." >&2
    echo "Remediation: Please ensure src/07_ssm_parameter_parser.sh is present in the script directory or in /usr/local/bin/src/." >&2
    echo "You may need to run 'make install' or consult the installation instructions in the README to properly set up the required modules." >&2
    exit 1
fi

# Source version module
if [ -f "${SCRIPT_DIR}/src/00_version.sh" ]; then
    # shellcheck source=./src/00_version.sh
    source "${SCRIPT_DIR}/src/00_version.sh"
elif [ -f "/usr/local/bin/src/00_version.sh" ]; then
    # shellcheck source=/dev/null
    source "/usr/local/bin/src/00_version.sh"
else
    # Set fallback version if module not found
    ZTIAWS_VERSION="unknown"
fi

# Function to show usage
usage() {
    cat << EOF
AWS SSM Session Manager Helper

POSITIONAL SYNTAX (current - backward compatible):
  $(basename "$0") [region] [instance-id]    # Connect to instance
  $(basename "$0") [instance-id]             # Uses default region (Canada Central)
  $(basename "$0") [region]                  # List instances in region

FLAG-BASED SYNTAX (new - enterprise friendly):
  $(basename "$0") --region [region] --instance [instance]  # Connect to instance
  $(basename "$0") --region [region] --list                # List instances
  $(basename "$0") --region [region]                       # List instances

MIXED SYNTAX (also supported):
  $(basename "$0") cac1 --instance i-1234567890abcd        # Mixed positional/flag
  $(basename "$0") --region cac1 i-1234567890abcd          # Mixed flag/positional

Supported regions:
  cac1  - Canada Central (Montreal) [default]
  apse1 - Asia Pacific (Seoul)
  caw1  - Canada West (Calgary)
  usw1  - US West (N. California)
  use1  - US East (N. Virginia)
  euw1  - EU West (Ireland)

Commands:
  $(basename "$0") install              # Show installation instructions
  $(basename "$0") version              # Show version information
  $(basename "$0") check                # Check system requirements
  $(basename "$0") exec [options]        # Run a command on an instance
  $(basename "$0") exec-tagged [options] # Run a command on instances by tag
EOF

    if [ "$FILE_TRANSFER_AVAILABLE" = true ]; then
        cat << EOF
  $(basename "$0") upload [options]      # Upload a file to an instance
  $(basename "$0") download [options]    # Download a file from an instance
EOF
    fi

    cat << EOF

Run Command Options:
  $(basename "$0") exec <region> <instance-id> "<command>"
  $(basename "$0") exec-tagged <region> <tag-key> <tag-value> "<command>"
EOF

    if [ "$FILE_TRANSFER_AVAILABLE" = true ]; then
        cat << EOF

File Transfer Options:
  $(basename "$0") upload <region> <instance-id> <local-file> <remote-path>
  $(basename "$0") download <region> <instance-id> <remote-file> <local-path>
EOF
    fi

    cat << EOF

POSITIONAL EXAMPLES (current syntax):
  $(basename "$0") cac1                             # List instances in Canada Central
  $(basename "$0") use1 i-1234                      # Connect to instance in US East
  $(basename "$0") i-1234                           # Connect to instance in Canada Central
  $(basename "$0") exec cac1 i-1234 "systemctl status nginx"   # Run command on instance
  $(basename "$0") exec cac1 web-server-prod "systemctl status nginx"  # Run command using instance name
  $(basename "$0") exec-tagged use1 Role web "df -h" # Run command on instances with tag Role=web

FLAG-BASED EXAMPLES (new syntax):
  $(basename "$0") --region cac1 --list            # List instances in Canada Central
  $(basename "$0") --region use1 --instance i-1234 --connect  # Connect to instance
  $(basename "$0") --exec --region cac1 --instance i-1234 --command "systemctl status nginx"  # Run command
  $(basename "$0") --exec-tagged --region use1 --tag-key Role --tag-value web --command "df -h"  # Run on tagged instances

MIXED EXAMPLES (combined syntax):
  $(basename "$0") cac1 --instance i-1234          # Region positional, instance flag
  $(basename "$0") --region use1 i-1234            # Region flag, instance positional
EOF

    if [ "$FILE_TRANSFER_AVAILABLE" = true ]; then
        cat << EOF
  $(basename "$0") upload cac1 i-1234 ./config.txt /etc/app/config.txt    # Upload file
  $(basename "$0") download cac1 i-1234 /var/log/app.log ./app.log        # Download file

File Transfer Notes:
  - Files < 1MB: Direct transfer via SSM (faster)
  - Files ≥ 1MB: Transfer via S3 intermediary (reliable for large files)
  - S3 bucket is auto-created and files are auto-deleted after 24 hours
  - Instance must have AWS CLI installed and proper IAM permissions for S3 access (large files only)
EOF
    fi
}

# Version information
VERSION="${ZTIAWS_VERSION:-unknown}"
show_version() {
    echo "ztiaws/ssm v$VERSION - AWS SSM Session Management"
    if [ "$FILE_TRANSFER_AVAILABLE" = true ]; then
        echo "Features: Sessions, Commands, File Transfer"
    else
        echo "Features: Sessions, Commands"
    fi
}

# Detect OS and package manager
detect_os() {
    if [ -f /etc/os-release ]; then
        # shellcheck source=/dev/null
        . /etc/os-release
        OS=$NAME
        VERSION=$VERSION_ID
    else
        OS=$(uname -s)
        VERSION=$(uname -r)
    fi
}

# Check system requirements
check_requirements() {
    local missing_deps=0

    # Check AWS CLI
    if ! command -v aws >/dev/null 2>&1; then
        log_error "AWS CLI is not installed"
        log_info "Please install AWS CLI by following the official documentation:"
        log_info "https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html"
        log_info "Choose your platform and follow the installation steps provided"
        missing_deps=1
    fi

    # Check SSM plugin
    if ! command -v session-manager-plugin >/dev/null 2>&1; then
        log_error "AWS Session Manager plugin is not installed"
        log_info "Would you like to auto-install it? (yes/no): "
        read -r response
        if [[ "$response" =~ ^[Yy](es)?$ ]]; then
            install_ssm_plugin
        else
            log_info "Please install the Session Manager plugin by following the official documentation:"
            log_info "https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-install-plugin.html"
            missing_deps=1
        fi
    fi

    # Check jq
    if ! command -v jq >/dev/null 2>&1; then
        log_error "jq is not installed"
        log_info "jq is required for robust parsing of AWS command output."
        log_info "Please install jq. Examples:"
        log_info "  Ubuntu/Debian: sudo apt-get install jq"
        log_info "  macOS: brew install jq"
        log_info "  Other: https://stedolan.github.io/jq/download/"
        missing_deps=1
    fi

    # Check AWS credentials
    if ! aws sts get-caller-identity >/dev/null 2>&1; then
        log_error "AWS credentials are not configured or invalid"
        log_info "Run 'aws configure' to set up your credentials"
        missing_deps=1
    fi

    if [ $missing_deps -eq 0 ]; then
        log_info "All requirements met!"
    else
        return 1
    fi
}

# Install SSM plugin based on OS
install_ssm_plugin() {
    detect_os
    log_info "Installing SSM plugin for $OS..."
    
    case $OS in
        *"Red Hat"*|*"Fedora"*)
            sudo curl "https://s3.amazonaws.com/session-manager-downloads/plugin/latest/linux_64bit/session-manager-plugin" -o "/usr/local/bin/session-manager-plugin"
            sudo chmod +x /usr/local/bin/session-manager-plugin
            ;;
        *"Ubuntu"*|*"Debian"*)
            curl "https://s3.amazonaws.com/session-manager-downloads/plugin/latest/ubuntu_64bit/session-manager-plugin.deb" -o "session-manager-plugin.deb"
            sudo dpkg -i session-manager-plugin.deb
            rm session-manager-plugin.deb
            ;;
        *)
            log_error "Unsupported OS: $OS"
            log_info "Please visit: https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-install-plugin.html"
            exit 1
            ;;
    esac
    
    if command -v session-manager-plugin >/dev/null 2>&1; then
        log_info "SSM plugin installed successfully!"
    else
        log_error "Failed to install SSM plugin"
        exit 1
    fi
}

# Default region is now handled by the parameter parser

# Validate instance ID format
validate_instance_identifier() {
    local identifier="$1"
    
    # If it looks like an instance ID, validate the format
    if [[ "$identifier" =~ ^i-[a-zA-Z0-9]{8,}$ ]]; then
        return 0
    fi
    
    # If it could be an instance name, allow it
    if is_potential_instance_name "$identifier"; then
        return 0
    fi
    
    log_error "Invalid instance identifier format. Should be instance ID (i-xxxxxxxx) or instance name"
    return 1
}
# Exec command handler
handle_exec_command() {
    if [ $# -lt 3 ]; then
        log_error "Missing required parameters for exec command"
        echo "Usage: $(basename "$0") exec <region> <instance-identifier> \"<command>\""
        exit 1
    fi

    local region="$1"  # Already validated and translated by parameter parser
    local instance_identifier="$2"
    local command="$3"

    if ! validate_instance_identifier "$instance_identifier"; then
        exit 1
    fi

    log_info "Resolving instance identifier: $instance_identifier"
    if ! instance_id=$(resolve_instance_identifier "$instance_identifier" "$region"); then
    
        exit 1
    fi

    log_info "Executing command on instance: $instance_id ($instance_identifier)"
    
    run_remote_command "$instance_id" "$region" "$command"
    return $?
}

# Exec-tagged command handler
handle_exec_tagged_command() {
    if [ $# -lt 4 ]; then
        log_error "Missing required parameters for exec-tagged command"
        echo "Usage: $(basename "$0") exec-tagged <region> <tag-key> <tag-value> \"<command>\""
        exit 1
    fi

    local region="$1"  # Already validated and translated by parameter parser
    local tag_key="$2"
    local tag_value="$3"
    local command="$4"

    log_info "Executing command on instances with tag $tag_key=$tag_value in region $region"
    
    run_remote_command_tagged "$tag_key" "$tag_value" "$region" "$command"
    return $?
}

# Upload command handler
handle_upload_command() {
    if [ "$FILE_TRANSFER_AVAILABLE" != true ]; then
        log_error "File transfer module not available"
        log_error "Please ensure src/04_ssm_file_transfer.sh is present and properly sourced"
        exit 1
    fi

    if [ $# -lt 4 ]; then
        log_error "Missing required parameters for upload command"
        echo "Usage: $(basename "$0") upload <region> <instance-identifier> <local-file> <remote-path>"
        exit 1
    fi

    local region="$1"  # Already validated and translated by parameter parser
    local instance_identifier="$2"
    local local_file="$3"
    local remote_path="$4"

    if ! validate_instance_identifier "$instance_identifier"; then
        exit 1
    fi

    log_info "Uploading file: $local_file -> $remote_path"
    log_info "Target: $instance_identifier in $region"
    
    upload_file "$region" "$instance_identifier" "$local_file" "$remote_path"
    return $?
}

# Download command handler
handle_download_command() {
    if [ "$FILE_TRANSFER_AVAILABLE" != true ]; then
        log_error "File transfer module not available"
        log_error "Please ensure src/04_ssm_file_transfer.sh is present and properly sourced"
        exit 1
    fi

    if [ $# -lt 4 ]; then
        log_error "Missing required parameters for download command"
        echo "Usage: $(basename "$0") download <region> <instance-identifier> <remote-file> <local-path>"
        exit 1
    fi

    local region="$1"  # Already validated and translated by parameter parser
    local instance_identifier="$2"
    local remote_file="$3"
    local local_path="$4"

    if ! validate_instance_identifier "$instance_identifier"; then
        exit 1
    fi

    log_info "Downloading file: $remote_file -> $local_path"
    log_info "Source: $instance_identifier in $region"
    
    download_file "$region" "$instance_identifier" "$remote_file" "$local_path"
    return $?
}

# Main logic
main() {
    # Parse parameters using the new parameter parser
    if ! parse_ssm_parameters "$@"; then
        exit 1
    fi

    # Handle commands based on parsed parameters
    local operation
    local region
    local instance
    local command
    local tag_key
    local tag_value
    local local_file
    local remote_file
    local local_path
    local remote_path
    local local_port
    local remote_port
    
    operation="$(get_ssm_operation)"
    region="$(get_ssm_region)"
    instance="$(get_ssm_instance)"
    command="$(get_ssm_command)"
    tag_key="$(get_ssm_tag_key)"
    tag_value="$(get_ssm_tag_value)"
    local_file="$(get_ssm_local_file)"
    remote_file="$(get_ssm_remote_file)"
    local_path="$(get_ssm_local_path)"
    remote_path="$(get_ssm_remote_path)"
    local_port="$(get_ssm_local_port)"
    remote_port="$(get_ssm_remote_port)"

    # Handle system commands first
    if [[ "$(get_ssm_help)" == "true" ]]; then
        usage
        exit 0
    fi

    if [[ "$(get_ssm_version)" == "true" ]]; then
        show_version
        exit 0
    fi

    if [[ "$(get_ssm_check)" == "true" ]]; then
        check_requirements
        exit 0
    fi

    if [[ "$operation" == "install" ]]; then
        # Since we have interactive SSM plugin installation, just show AWS CLI instructions
        log_info "Installation Instructions:"
        echo
        log_info "1. AWS CLI Installation:"
        echo "   Visit: https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html"
        echo "   Follow the instructions for your platform"
        echo
        log_info "2. Session Manager Plugin:"
        echo "   Run '$(basename "$0") check' to install the plugin interactively"
        echo "   Or visit: https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-install-plugin.html"
        if [ "$FILE_TRANSFER_AVAILABLE" = true ]; then
            echo
            log_info "3. File Transfer Requirements:"
            echo "   - For large files (≥1MB): Instances need AWS CLI and S3 access"
            echo "   - For small files (<1MB): No additional requirements"
        fi
        exit 0
    fi

    # Check requirements before proceeding with actual commands
    if ! check_requirements >/dev/null 2>&1; then
        log_error "System requirements not met. Run '$(basename "$0") check' for details"
        exit 1
    fi

    # Execute operations based on parsed parameters
    case "$operation" in
        "list")
            log_info "Listing instances in $region..."
            aws ec2 describe-instances --region "$region" \
                --query "Reservations[*].Instances[*].{
                    Name: Tags[?Key=='Name'].Value | [0],
                    InstanceId: InstanceId,
                    IP: PrivateIpAddress,
                    State: State.Name,
                    OS: PlatformDetails
                }" \
                --output table
            echo -e "\nUsage: $(basename "$0") $region <instance-id>"
            ;;
        "connect")
            if ! validate_instance_identifier "$instance"; then
                exit 1
            fi
            
            # Resolve instance name to ID if needed
            local instance_id="$instance"
            if ! [[ "$instance" =~ ^i-[a-zA-Z0-9]{8,}$ ]]; then
                instance_id=$(resolve_instance_identifier "$instance" "$region")
                if ! [[ "$instance_id" =~ ^i-[a-zA-Z0-9]{8,}$ ]]; then
                    exit 1
                fi
            fi
            
            log_info "Connecting to instance $instance_id in $region..."
            aws ssm start-session \
                --region "$region" \
                --target "$instance_id"
            ;;
        "exec")
            handle_exec_command "$region" "$instance" "$command"
            exit $?
            ;;
        "exec-tagged")
            handle_exec_tagged_command "$region" "$tag_key" "$tag_value" "$command"
            exit $?
            ;;
        "upload")
            if [ "$FILE_TRANSFER_AVAILABLE" != true ]; then
                log_error "File transfer module not available"
                log_error "Please ensure src/04_ssm_file_transfer.sh is present and properly sourced"
                exit 1
            fi
            handle_upload_command "$region" "$instance" "$local_file" "$remote_path"
            exit $?
            ;;
        "download")
            if [ "$FILE_TRANSFER_AVAILABLE" != true ]; then
                log_error "File transfer module not available"
                log_error "Please ensure src/04_ssm_file_transfer.sh is present and properly sourced"
                exit 1
            fi
            handle_download_command "$region" "$instance" "$remote_file" "$local_path"
            exit $?
            ;;
        "forward")
            # Resolve instance identifier to instance ID if necessary
            local instance_id="$instance"
            if ! [[ "$instance" =~ ^i-[a-zA-Z0-9]{8,}$ ]]; then
                instance_id=$(resolve_instance_identifier "$instance" "$region")
                if ! [[ "$instance_id" =~ ^i-[a-zA-Z0-9]{8,}$ ]]; then
                    exit 1
                fi
            fi
            
            log_info "Setting up port forwarding: localhost:$local_port -> $instance_id:$remote_port"
            aws ssm start-session \
                --region "$region" \
                --target "$instance_id" \
                --document-name "AWS-StartPortForwardingSession" \
                --parameters "portNumber=$remote_port,localPortNumber=$local_port"
            ;;
        "")
            # No operation determined, show usage
            usage
            exit 0
            ;;
        *)
            log_error "Unknown operation: $operation"
            usage
            exit 1
            ;;
    esac
}

# Run main function
main "$@"