#!/bin/bash
# Security test script for restricted shell
# Tests various injection and escape attempts

set -e

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

echo -e "${BLUE}╔════════════════════════════════════════════════════════╗${NC}"
echo -e "${BLUE}║   Restricted Shell Security Test                      ║${NC}"
echo -e "${BLUE}╚════════════════════════════════════════════════════════╝${NC}"
echo ""

PASSED=0
FAILED=0
TOTAL=0

# Function to test a command
test_command() {
    local test_name="$1"
    local test_input="$2"
    local should_block="$3"  # "block" or "allow"

    TOTAL=$((TOTAL + 1))

    echo -ne "${YELLOW}Testing:${NC} $test_name ... "

    # Note: This is a template. In practice, you'd need to:
    # 1. Send input to the restricted shell
    # 2. Check if it was blocked or executed
    # 3. Verify no unauthorized commands ran

    # For now, we'll test the regex patterns
    if [[ "$should_block" == "block" ]]; then
        # These should be blocked
        if [[ "$test_input" =~ ^-P[[:space:]]+[a-zA-Z0-9_-]+$ ]] || \
           [[ "$test_input" =~ ^wss?://[a-zA-Z0-9\.\:/_-]+$ ]]; then
            echo -e "${RED}FAIL${NC} - Should have blocked but pattern matched"
            FAILED=$((FAILED + 1))
        else
            echo -e "${GREEN}PASS${NC} - Correctly blocked"
            PASSED=$((PASSED + 1))
        fi
    else
        # These should be allowed
        if [[ "$test_input" =~ ^-P[[:space:]]+[a-zA-Z0-9_-]+$ ]] || \
           [[ "$test_input" =~ ^wss?://[a-zA-Z0-9\.\:/_-]+$ ]]; then
            echo -e "${GREEN}PASS${NC} - Correctly allowed"
            PASSED=$((PASSED + 1))
        else
            echo -e "${RED}FAIL${NC} - Should have allowed but pattern didn't match"
            FAILED=$((FAILED + 1))
        fi
    fi
}

echo -e "${BLUE}═══ Testing Valid Commands (Should Allow) ═══${NC}"
echo ""

test_command "Local profile" "-P local" "allow"
test_command "Remote profile" "-P rpi-master" "allow"
test_command "Profile with dash" "-P rpi-worker-1" "allow"
test_command "Profile with underscore" "-P my_profile" "allow"
test_command "Websocket URL" "ws://192.168.1.100:3000" "allow"
test_command "Secure websocket" "wss://example.com:3000" "allow"
test_command "Websocket with path" "ws://192.168.1.100:3000/ws" "allow"

echo ""
echo -e "${BLUE}═══ Testing Command Injection (Should Block) ═══${NC}"
echo ""

test_command "Command substitution \$()" "-P \$(whoami)" "block"
test_command "Command substitution backticks" "-P \`id\`" "block"
test_command "Shell semicolon" "-P local; ls -la" "block"
test_command "Shell AND operator" "-P local && cat /etc/passwd" "block"
test_command "Shell OR operator" "-P local || /bin/sh" "block"
test_command "Shell pipe" "-P local | grep root" "block"
test_command "Shell redirect" "-P local > /tmp/output" "block"
test_command "Shell background" "-P local &" "block"

echo ""
echo -e "${BLUE}═══ Testing Path Traversal (Should Block) ═══${NC}"
echo ""

test_command "Parent directory" "-P ../etc/passwd" "block"
test_command "Absolute path" "-P /etc/passwd" "block"
test_command "Multiple parent dirs" "-P ../../bin/bash" "block"
test_command "Encoded path" "-P %2e%2e%2f" "block"

echo ""
echo -e "${BLUE}═══ Testing Special Characters (Should Block) ═══${NC}"
echo ""

test_command "Newline injection" "-P local\nls" "block"
test_command "Carriage return" "-P local\rls" "block"
test_command "Null byte" "-P local\x00ls" "block"
test_command "Single quote" "-P local' ls" "block"
test_command "Double quote" "-P local\" ls" "block"
test_command "Dollar sign" "-P \$HOME" "block"
test_command "Asterisk wildcard" "-P local*" "block"
test_command "Question wildcard" "-P local?" "block"

echo ""
echo -e "${BLUE}═══ Testing Environment Variables (Should Block) ═══${NC}"
echo ""

test_command "HOME variable" "-P \$HOME" "block"
test_command "PATH variable" "-P \$PATH" "block"
test_command "SHELL variable" "-P \$SHELL" "block"
test_command "Braced variable" "-P \${HOME}" "block"

echo ""
echo -e "${BLUE}═══ Testing WebSocket URL Exploits (Should Block) ═══${NC}"
echo ""

test_command "WS with command injection" "ws://evil.com/\$(id)" "block"
test_command "WS with backticks" "ws://evil.com/\`whoami\`" "block"
test_command "WS with semicolon" "ws://evil.com/; ls" "block"
test_command "WS with spaces" "ws://evil.com/ /bin/sh" "block"

echo ""
echo -e "${BLUE}════════════════════════════════════════════════════${NC}"
echo -e "${BLUE}                  TEST SUMMARY                       ${NC}"
echo -e "${BLUE}════════════════════════════════════════════════════${NC}"
echo ""
echo -e "Total Tests:  ${BLUE}$TOTAL${NC}"
echo -e "Passed:       ${GREEN}$PASSED${NC}"
echo -e "Failed:       ${RED}$FAILED${NC}"
echo ""

if [ $FAILED -eq 0 ]; then
    echo -e "${GREEN}✓ All security tests passed!${NC}"
    echo ""
    echo -e "${YELLOW}Note:${NC} These are pattern validation tests only."
    echo "For complete security verification, you should:"
    echo "  1. Test in actual container environment"
    echo "  2. Verify socktop binary doesn't process malicious args"
    echo "  3. Monitor for unexpected process execution"
    echo "  4. Check logs for injection attempts"
    echo ""
    exit 0
else
    echo -e "${RED}✗ Some security tests failed!${NC}"
    echo ""
    echo "Review the failed tests and update regex patterns in restricted-shell.sh"
    echo ""
    exit 1
fi