#!/bin/bash
# SSH Performance Benchmark
# Tests: rsync, SCP, exec commands

set -e

SSH_HOST="127.0.0.1"
SSH_PORT="2024"
SSH_USER="demo"
SSH_PASS="demo123"
SSH_KEY="/Users/accusys/.ssh/id_rsa"
TEST_DIR="/tmp/ssh_benchmark"
REMOTE_DIR="/Users/accusys/momentry/var/sftpgo/data/demo/benchmark"

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

echo "================================================"
echo "SSH Performance Benchmark"
echo "================================================"
echo ""

# Setup test directory
rm -rf "$TEST_DIR"
mkdir -p "$TEST_DIR"
cd "$TEST_DIR"

# Generate test files
echo "Generating test files..."
dd if=/dev/urandom of=file_1mb.bin bs=1M count=1 2>/dev/null
dd if=/dev/urandom of=file_10mb.bin bs=1M count=10 2>/dev/null
dd if=/dev/urandom of=file_50mb.bin bs=1M count=50 2>/dev/null
dd if=/dev/urandom of=file_100mb.bin bs=1M count=100 2>/dev/null
echo "${GREEN}Test files generated${NC}"
echo ""

# Setup SSH key auth (if needed)
# sshpass -p "$SSH_PASS" ssh-copy-id -p "$SSH_PORT" "$SSH_USER@$SSH_HOST" 2>/dev/null || true

# Test 1: SCP upload 1MB
echo "=== Test 1: SCP upload 1MB ==="
START=$(date +%s.%N)
sshpass -p "$SSH_PASS" scp -P "$SSH_PORT" file_1mb.bin "$SSH_USER@$SSH_HOST:$REMOTE_DIR/file_1mb.bin"
END=$(date +%s.%N)
ELAPSED=$(echo "$END - $START" | bc)
THROUGHPUT=$(echo "scale=2; 1 / $ELAPSED" | bc)
echo "${GREEN}SCP upload 1MB: ${ELAPSED}s (${THROUGHPUT} MB/s)${NC}"
echo ""

# Test 2: SCP upload 10MB
echo "=== Test 2: SCP upload 10MB ==="
START=$(date +%s.%N)
sshpass -p "$SSH_PASS" scp -P "$SSH_PORT" file_10mb.bin "$SSH_USER@$SSH_HOST:$REMOTE_DIR/file_10mb.bin"
END=$(date +%s.%N)
ELAPSED=$(echo "$END - $START" | bc)
THROUGHPUT=$(echo "scale=2; 10 / $ELAPSED" | bc)
echo "${GREEN}SCP upload 10MB: ${ELAPSED}s (${THROUGHPUT} MB/s)${NC}"
echo ""

# Test 3: SCP download 10MB
echo "=== Test 3: SCP download 10MB ==="
START=$(date +%s.%N)
sshpass -p "$SSH_PASS" scp -P "$SSH_PORT" "$SSH_USER@$SSH_HOST:$REMOTE_DIR/file_10mb.bin" /tmp/scp_download_10mb.bin
END=$(date +%s.%N)
ELAPSED=$(echo "$END - $START" | bc)
THROUGHPUT=$(echo "scale=2; 10 / $ELAPSED" | bc)
echo "${GREEN}SCP download 10MB: ${ELAPSED}s (${THROUGHPUT} MB/s)${NC}"
echo ""

# Test 4: rsync upload 50MB
echo "=== Test 4: rsync upload 50MB ==="
START=$(date +%s.%N)
sshpass -p "$SSH_PASS" rsync -avz -e "ssh -p $SSH_PORT" file_50mb.bin "$SSH_USER@$SSH_HOST:$REMOTE_DIR/"
END=$(date +%s.%N)
ELAPSED=$(echo "$END - $START" | bc)
THROUGHPUT=$(echo "scale=2; 50 / $ELAPSED" | bc)
echo "${GREEN}rsync upload 50MB: ${ELAPSED}s (${THROUGHPUT} MB/s)${NC}"
echo ""

# Test 5: rsync download 50MB
echo "=== Test 5: rsync download 50MB ==="
START=$(date +%s.%N)
sshpass -p "$SSH_PASS" rsync -avz -e "ssh -p $SSH_PORT" "$SSH_USER@$SSH_HOST:$REMOTE_DIR/file_50mb.bin" /tmp/
END=$(date +%s.%N)
ELAPSED=$(echo "$END - $START" | bc)
THROUGHPUT=$(echo "scale=2; 50 / $ELAPSED" | bc)
echo "${GREEN}rsync download 50MB: ${ELAPSED}s (${THROUGHPUT} MB/s)${NC}"
echo ""

# Test 6: rsync upload 100MB
echo "=== Test 6: rsync upload 100MB ==="
START=$(date +%s.%N)
sshpass -p "$SSH_PASS" rsync -avz -e "ssh -p $SSH_PORT" file_100mb.bin "$SSH_USER@$SSH_HOST:$REMOTE_DIR/"
END=$(date +%s.%N)
ELAPSED=$(echo "$END - $START" | bc)
THROUGHPUT=$(echo "scale=2; 100 / $ELAPSED" | bc)
echo "${GREEN}rsync upload 100MB: ${ELAPSED}s (${THROUGHPUT} MB/s)${NC}"
echo ""

# Test 7: rsync download 100MB
echo "=== Test 7: rsync download 100MB ==="
START=$(date +%s.%N)
sshpass -p "$SSH_PASS" rsync -avz -e "ssh -p $SSH_PORT" "$SSH_USER@$SSH_HOST:$REMOTE_DIR/file_100mb.bin" /tmp/
END=$(date +%s.%N)
ELAPSED=$(echo "$END - $START" | bc)
THROUGHPUT=$(echo "scale=2; 100 / $ELAPSED" | bc)
echo "${GREEN}rsync download 100MB: ${ELAPSED}s (${THROUGHPUT} MB/s)${NC}"
echo ""

# Test 8: rsync delta transfer (modify 10MB and re-upload)
echo "=== Test 8: rsync delta transfer ==="
dd if=/dev/urandom of=file_10mb.bin bs=1M count=10 2>/dev/null
# Modify first 1MB only
dd if=/dev/urandom of=file_10mb.bin bs=1M count=1 seek=0 conv=notrunc 2>/dev/null
START=$(date +%s.%N)
sshpass -p "$SSH_PASS" rsync -avz -e "ssh -p $SSH_PORT" file_10mb.bin "$SSH_USER@$SSH_HOST:$REMOTE_DIR/"
END=$(date +%s.%N)
ELAPSED=$(echo "$END - $START" | bc)
echo "${GREEN}rsync delta (1MB change): ${ELAPSED}s${NC}"
echo ""

# Test 9: exec command latency
echo "=== Test 9: exec command latency ==="
START=$(date +%s.%N)
sshpass -p "$SSH_PASS" ssh -p "$SSH_PORT" "$SSH_USER@$SSH_HOST" "echo test" > /dev/null
END=$(date +%s.%N)
ELAPSED=$(echo "$END - $START" | bc)
echo "${GREEN}exec 'echo test': ${ELAPSED}s${NC}"
echo ""

# Test 10: exec command throughput (pipe 10MB)
echo "=== Test 10: exec command throughput (pipe 10MB) ==="
START=$(date +%s.%N)
sshpass -p "$SSH_PASS" ssh -p "$SSH_PORT" "$SSH_USER@$SSH_HOST" "cat > $REMOTE_DIR/pipe_10mb.bin" < file_10mb.bin
END=$(date +%s.%N)
ELAPSED=$(echo "$END - $START" | bc)
THROUGHPUT=$(echo "scale=2; 10 / $ELAPSED" | bc)
echo "${GREEN}exec pipe 10MB: ${ELAPSED}s (${THROUGHPUT} MB/s)${NC}"
echo ""

# Cleanup remote files
echo "=== Cleanup ==="
sshpass -p "$SSH_PASS" ssh -p "$SSH_PORT" "$SSH_USER@$SSH_HOST" "rm -rf $REMOTE_DIR/*" 2>/dev/null || true

# Cleanup local
rm -rf "$TEST_DIR"
rm -f /tmp/scp_download_*.bin /tmp/file_*.bin

echo "================================================"
echo "${GREEN}SSH Performance Benchmark Complete${NC}"
echo "================================================"