#!/bin/bash
#
# Citadel Network Policy Test Script
# Verifies zero-trust policies are working correctly
#

set -e

echo "=========================================="
echo "  Citadel Network Policy Verification"
echo "  Testing Zero-Trust Security"
echo "=========================================="
echo ""

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

# Set kubeconfig
export KUBECONFIG=/etc/rancher/k3s/k3s.yaml

PASSED=0
FAILED=0

test_pass() {
    echo -e "  ${GREEN}✓ PASS${NC}: $1"
    ((PASSED++))
}

test_fail() {
    echo -e "  ${RED}✗ FAIL${NC}: $1"
    ((FAILED++))
}

test_connectivity() {
    local from_pod=$1
    local to_service=$2
    local port=$3
    local should_work=$4
    local description=$5

    echo -e "${BLUE}Testing:${NC} $description"

    # Try to connect
    if kubectl exec -n demo "$from_pod" -- timeout 3 sh -c "echo '' | nc -w 2 $to_service $port" &>/dev/null; then
        if [ "$should_work" = "true" ]; then
            test_pass "Connection succeeded (expected)"
        else
            test_fail "Connection succeeded (should have been blocked!)"
        fi
    else
        if [ "$should_work" = "false" ]; then
            test_pass "Connection blocked (expected)"
        else
            test_fail "Connection blocked (should have succeeded!)"
        fi
    fi
}

# Wait for test-client pod
echo -e "${YELLOW}Waiting for test-client pod...${NC}"
kubectl wait --for=condition=Ready pod/test-client -n demo --timeout=60s 2>/dev/null || {
    echo -e "${RED}test-client pod not ready. Deploy demo app first.${NC}"
    exit 1
}
echo ""

echo "=========================================="
echo "  Test Suite: Allowed Connections"
echo "=========================================="
echo ""

# Test 1: test-client -> frontend (should work)
test_connectivity "test-client" "frontend" "80" "true" \
    "test-client -> frontend:80"

# Test 2: test-client -> backend (should work)
test_connectivity "test-client" "backend" "5678" "true" \
    "test-client -> backend:5678"

# Get a frontend pod
FRONTEND_POD=$(kubectl get pods -n demo -l tier=frontend -o jsonpath='{.items[0].metadata.name}')

# Test 3: frontend -> backend (should work)
echo -e "${BLUE}Testing:${NC} frontend -> backend:5678"
if kubectl exec -n demo "$FRONTEND_POD" -- timeout 3 sh -c "wget -q -O- backend:5678" &>/dev/null; then
    test_pass "Connection succeeded (expected)"
else
    test_fail "Connection blocked (should have succeeded!)"
fi

# Get a backend pod
BACKEND_POD=$(kubectl get pods -n demo -l tier=backend -o jsonpath='{.items[0].metadata.name}')

# Test 4: backend -> database (should work)
echo -e "${BLUE}Testing:${NC} backend -> database:6379"
if kubectl exec -n demo "$BACKEND_POD" -- timeout 3 sh -c "echo PING | nc -w 2 database 6379" 2>/dev/null | grep -q "PONG"; then
    test_pass "Connection succeeded (expected)"
else
    # Try alternate method
    if kubectl exec -n demo "$BACKEND_POD" -- timeout 3 sh -c "echo '' | nc -w 2 database 6379" &>/dev/null; then
        test_pass "Connection succeeded (expected)"
    else
        test_fail "Connection blocked (should have succeeded!)"
    fi
fi

echo ""
echo "=========================================="
echo "  Test Suite: Blocked Connections"
echo "=========================================="
echo ""

# Test 5: test-client -> database (should be blocked)
test_connectivity "test-client" "database" "6379" "false" \
    "test-client -> database:6379 (should be blocked)"

# Test 6: frontend -> database (should be blocked)
echo -e "${BLUE}Testing:${NC} frontend -> database:6379 (should be blocked)"
if kubectl exec -n demo "$FRONTEND_POD" -- timeout 3 sh -c "echo '' | nc -w 2 database 6379" &>/dev/null; then
    test_fail "Connection succeeded (should have been blocked!)"
else
    test_pass "Connection blocked (expected)"
fi

# Test 7: database -> backend (egress blocked)
DATABASE_POD=$(kubectl get pods -n demo -l tier=database -o jsonpath='{.items[0].metadata.name}')
echo -e "${BLUE}Testing:${NC} database -> backend:5678 (egress should be blocked)"
if kubectl exec -n demo "$DATABASE_POD" -- timeout 3 sh -c "echo '' | nc -w 2 backend 5678" &>/dev/null; then
    test_fail "Connection succeeded (should have been blocked!)"
else
    test_pass "Connection blocked (expected)"
fi

echo ""
echo "=========================================="
echo "  Test Results Summary"
echo "=========================================="
echo ""
echo -e "  ${GREEN}Passed:${NC} $PASSED"
echo -e "  ${RED}Failed:${NC} $FAILED"
echo ""

if [ $FAILED -eq 0 ]; then
    echo -e "${GREEN}All tests passed! Zero-trust policies working correctly.${NC}"
    exit 0
else
    echo -e "${RED}Some tests failed. Review network policies.${NC}"
    echo ""
    echo "Debug commands:"
    echo "  hubble observe --namespace demo --verdict DROPPED"
    echo "  kubectl get ciliumnetworkpolicies -n demo -o yaml"
    exit 1
fi
