#!/bin/bash
set -euo pipefail

# nbcrypt - SSH Agent based encryption/decryption tool
# Uses Ed25519 signature for deterministic key derivation

SCRIPT_NAME="$(basename "$0")"
KEY_SEED_MESSAGE="nbase2-secret-key-seed-v1"
KEY_IDENTITY="id_ed25519"

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

error() {
    echo -e "${RED}❌ Error: $*${NC}" >&2
    exit 1
}

info() {
    echo -e "${GREEN}ℹ️  $*${NC}" >&2
}

warn() {
    echo -e "${YELLOW}⚠️  $*${NC}" >&2
}

usage() {
    cat <<EOF
Usage: $SCRIPT_NAME <command> [arguments]

Commands:
    encrypt <input> <output>   Encrypt a file using SSH Agent key
    decrypt <input> <output>   Decrypt a file using SSH Agent key
    check                      Check if SSH Agent has required key
    help                       Show this help message

Requirements:
    - SSH Agent must be running with id_ed25519 key loaded
    - ssh-keygen and openssl commands must be available

Examples:
    $SCRIPT_NAME encrypt secrets.txt secrets.enc
    $SCRIPT_NAME decrypt secrets.enc secrets.txt
    $SCRIPT_NAME check

EOF
    exit 0
}

check_dependencies() {
    local missing=()
    
    for cmd in ssh-keygen openssl ssh-add; do
        if ! command -v "$cmd" >/dev/null 2>&1; then
            missing+=("$cmd")
        fi
    done
    
    if [ ${#missing[@]} -gt 0 ]; then
        error "Missing required commands: ${missing[*]}"
    fi
}

check_ssh_agent() {
    # Check if SSH_AUTH_SOCK is set
    if [ -z "${SSH_AUTH_SOCK:-}" ]; then
        error "SSH Agent is not running or SSH_AUTH_SOCK is not set"
    fi
    
    # Check if ssh-add can connect to agent
    if ! ssh-add -l >/dev/null 2>&1; then
        local rc=$?
        if [ $rc -eq 2 ]; then
            error "Cannot connect to SSH Agent"
        elif [ $rc -eq 1 ]; then
            error "SSH Agent has no identities loaded"
        fi
    fi
    
    # Check if Ed25519 key is loaded
    if ! ssh-add -l 2>/dev/null | grep -q "ED25519"; then
        error "No Ed25519 key found in SSH Agent. Please add id_ed25519 with: ssh-add ~/.ssh/id_ed25519"
    fi
    
    info "SSH Agent check passed (Ed25519 key found)"
    return 0
}

derive_key() {
    # Derive encryption key from Ed25519 signature
    # Ed25519 signatures are deterministic, so same message = same signature
    
    local temp_message=$(mktemp)
    local temp_sig=$(mktemp)
    local temp_pubkey=$(mktemp)
    
    trap "rm -f '$temp_message' '$temp_sig' '$temp_pubkey'" EXIT
    
    # Create message to sign
    echo -n "$KEY_SEED_MESSAGE" > "$temp_message"
    
    # Get public key from agent for the Ed25519 key
    ssh-add -L 2>/dev/null | grep "ssh-ed25519" | head -1 > "$temp_pubkey"
    
    if [ ! -s "$temp_pubkey" ]; then
        error "Could not extract Ed25519 public key from SSH Agent"
    fi
    
    # Sign the message using ssh-keygen
    # Note: ssh-keygen -Y sign requires the public key in "allowed signers" format
    local temp_allowed=$(mktemp)
    echo "key1 $(cat "$temp_pubkey")" > "$temp_allowed"
    
    # Sign the message
    if ! ssh-keygen -Y sign -f "$temp_allowed" -n "nbase2" < "$temp_message" > "$temp_sig" 2>/dev/null; then
        # If -Y sign doesn't work (older ssh-keygen), try alternative method
        # We'll use ssh-agent's signing capability directly through a workaround
        
        # Extract just the signature data using ssh-agent protocol
        # Since we can't easily do that in pure bash, we'll use a deterministic approach
        # based on the public key itself as a fallback
        
        warn "ssh-keygen -Y sign not available, using fallback key derivation"
        
        # Fallback: Use public key hash as seed
        cat "$temp_pubkey" | sha256sum | head -c 64
        rm -f "$temp_allowed"
        return 0
    fi
    
    rm -f "$temp_allowed"
    
    # Extract signature and hash it to create 256-bit key
    # The signature file contains base64-encoded signature data
    cat "$temp_sig" | base64 -d 2>/dev/null | sha256sum | head -c 64
}

encrypt_file() {
    local input="$1"
    local output="$2"
    
    if [ ! -f "$input" ]; then
        error "Input file not found: $input"
    fi
    
    check_ssh_agent
    
    info "Deriving encryption key from SSH Agent..."
    local key=$(derive_key)
    
    if [ -z "$key" ]; then
        error "Failed to derive encryption key"
    fi
    
    info "Encrypting $input -> $output..."
    
    # Use OpenSSL to encrypt with AES-256-CBC
    # Pass the hex key directly to openssl
    openssl enc -aes-256-cbc -salt -pbkdf2 -in "$input" -out "$output" -pass "pass:$key"
    
    if [ $? -eq 0 ]; then
        info "✅ Encryption successful: $output"
    else
        error "Encryption failed"
    fi
}

decrypt_file() {
    local input="$1"
    local output="$2"
    
    if [ ! -f "$input" ]; then
        error "Input file not found: $input"
    fi
    
    check_ssh_agent
    
    info "Deriving decryption key from SSH Agent..."
    local key=$(derive_key)
    
    if [ -z "$key" ]; then
        error "Failed to derive decryption key"
    fi
    
    info "Decrypting $input -> $output..."
    
    # Use OpenSSL to decrypt with AES-256-CBC
    if openssl enc -aes-256-cbc -d -salt -pbkdf2 -in "$input" -out "$output" -pass "pass:$key" 2>/dev/null; then
        info "✅ Decryption successful: $output"
    else
        error "Decryption failed (wrong key or corrupted file)"
    fi
}

# Main command dispatcher
main() {
    if [ $# -eq 0 ]; then
        usage
    fi
    
    check_dependencies
    
    local command="$1"
    shift
    
    case "$command" in
        encrypt)
            if [ $# -ne 2 ]; then
                error "encrypt requires 2 arguments: <input> <output>"
            fi
            encrypt_file "$1" "$2"
            ;;
        decrypt)
            if [ $# -ne 2 ]; then
                error "decrypt requires 2 arguments: <input> <output>"
            fi
            decrypt_file "$1" "$2"
            ;;
        check)
            check_ssh_agent
            echo "✅ All checks passed"
            ;;
        help|--help|-h)
            usage
            ;;
        *)
            error "Unknown command: $command (use 'help' for usage)"
            ;;
    esac
}

main "$@"