#!/usr/bin/env python3
"""
Comprehensive tests for Phase 2 Trigger System implementation.
Run these tests to verify all trigger types work correctly.
"""
import asyncio
import pytest
import time
from unittest.mock import Mock, AsyncMock, patch
from pentest_mcp_server.trigger_system import TriggerSystem, OutputBuffer
from pentest_mcp_server.ssh_manager import SSHManager
from pentest_mcp_server.tmux_manager import TmuxManager
class TestTriggerSystem:
"""Test cases for the enhanced trigger system."""
@pytest.fixture
def ssh_manager(self):
"""Create mock SSH manager."""
ssh_manager = Mock(spec=SSHManager)
ssh_manager.run_command = AsyncMock()
return ssh_manager
@pytest.fixture
def tmux_manager(self):
"""Create mock Tmux manager."""
tmux_manager = Mock(spec=TmuxManager)
tmux_manager.capture_pane = AsyncMock()
tmux_manager.execute_command = AsyncMock()
return tmux_manager
@pytest.fixture
def trigger_system(self, ssh_manager, tmux_manager):
"""Create TriggerSystem instance."""
return TriggerSystem(ssh_manager, tmux_manager)
def test_prompt_detection(self, trigger_system):
"""Test prompt detection with various shell prompts."""
test_cases = [
("root@kali:~# ", True),
("kali@localhost:$ ", True),
("msf6 > ", True),
("(meterpreter) > ", True),
("mysql> ", True),
(">>> ", True),
("In [1]: ", True),
("regular text output", False),
("[sudo] password for kali:", True),
]
for output, expected in test_cases:
result = trigger_system._check_prompt_trigger(output)
assert result["matched"] == expected, f"Failed for: {output}"
print("Prompt detection works for all shell types")
def test_regex_trigger(self, trigger_system):
"""Test regex pattern matching."""
trigger = {"type": "regex", "pattern": "Nmap scan report", "name": "scan_started"}
# Test match in new output
result = trigger_system._check_regex_trigger(
trigger,
"Starting scan...\nNmap scan report for 192.168.1.1",
"Full output"
)
assert result["matched"] == True
assert result["trigger_name"] == "scan_started"
# Test no match
result = trigger_system._check_regex_trigger(
trigger,
"Starting scan...",
"Full output"
)
assert result["matched"] == False
print("Regex trigger detection works")
def test_timeout_trigger(self, trigger_system):
"""Test timeout trigger logic."""
start_time = time.time() - 10 # 10 seconds ago
# Test timeout reached
trigger = {"type": "timeout", "timeout_seconds": 5}
result = trigger_system._check_timeout_trigger(start_time, trigger)
assert result["matched"] == True
assert result["elapsed_seconds"] >= 10
# Test timeout not reached
start_time = time.time() - 2 # 2 seconds ago
trigger = {"type": "timeout", "timeout_seconds": 5}
result = trigger_system._check_timeout_trigger(start_time, trigger)
assert result["matched"] == False
print("Timeout trigger works")
@pytest.mark.asyncio
async def test_file_exists_trigger(self, trigger_system, tmux_manager):
"""Test file existence checking."""
trigger = {"type": "file_exists", "path": "/tmp/test.txt", "name": "file_ready"}
# Mock file exists
tmux_manager.execute_command.return_value = {"status": "sent"}
tmux_manager.capture_pane.return_value = {
"status": "success",
"output": "EXISTS\n1234" # File exists with size 1234
}
result = await trigger_system._check_file_exists_trigger(trigger, "test_session")
assert result["matched"] == True
assert result["file_path"] == "/tmp/test.txt"
assert result["file_size"] == 1234
# Mock file doesn't exist
tmux_manager.capture_pane.return_value = {
"status": "success",
"output": "NOTFOUND"
}
result = await trigger_system._check_file_exists_trigger(trigger, "test_session")
assert result["matched"] == False
print("File exists trigger works")
def test_error_detection(self, trigger_system):
"""Test automatic error detection in output."""
error_cases = [
("Error: connection failed", True),
("Command not found: nmap", True),
("Permission denied", True),
("Connection refused", True),
("No route to host", True),
("Segmentation fault", True),
("Normal output here", False),
("This is fine", False),
]
for output, should_detect_error in error_cases:
result = trigger_system._detect_error_in_output(output)
assert result["error_detected"] == should_detect_error, f"Failed for: {output}"
print("Error detection works for common error patterns")
def test_output_buffer(self):
"""Test efficient output buffer management."""
buffer = OutputBuffer()
# First call - should return all output as new
new, full = buffer.get_new_output("line1\nline2\nline3")
assert new == "line1\nline2\nline3"
assert full == "line1\nline2\nline3"
# Second call with additional lines - should return only new lines
new, full = buffer.get_new_output("line1\nline2\nline3\nline4\nline5")
assert new == "line4\nline5"
assert full == "line1\nline2\nline3\nline4\nline5"
# Third call with same output - should return empty
new, full = buffer.get_new_output("line1\nline2\nline3\nline4\nline5")
assert new == ""
assert full == "line1\nline2\nline3\nline4\nline5"
print("Output buffer management works efficiently")
def test_trigger_validation(self, trigger_system):
"""Test trigger configuration validation."""
# Valid triggers
valid_triggers = [
[{"type": "prompt"}],
[{"type": "regex", "pattern": "test"}],
[{"type": "timeout", "timeout_seconds": 30}],
[{"type": "file_exists", "path": "/tmp/test"}],
]
for triggers in valid_triggers:
is_valid, error = trigger_system.validate_triggers(triggers)
assert is_valid == True, f"Should be valid: {triggers}"
# Invalid triggers
invalid_cases = [
([{"type": "regex"}], "pattern"),
([{"type": "timeout"}], "timeout_seconds"),
([{"type": "file_exists"}], "path"),
([{"type": "unknown_type"}], "unknown"),
]
for triggers, expected_error in invalid_cases:
is_valid, error = trigger_system.validate_triggers(triggers)
assert is_valid == False, f"Should be invalid: {triggers}"
assert expected_error in error.lower()
print("Trigger validation works correctly")
class TestRealWorldScenarios:
"""Test real-world pentesting scenarios with triggers."""
@pytest.fixture
def mock_managers(self):
"""Set up mock managers for integration tests."""
ssh_manager = Mock(spec=SSHManager)
tmux_manager = Mock(spec=TmuxManager)
# Mock successful command execution
tmux_manager.execute_command.return_value = {"status": "sent"}
return ssh_manager, tmux_manager
@pytest.mark.asyncio
async def test_nmap_scan_scenario(self, mock_managers):
"""Test nmap scan with completion trigger."""
ssh_manager, tmux_manager = mock_managers
# Simulate nmap output progression
nmap_outputs = [
"Starting Nmap 7.80...",
"Nmap scan report for 192.168.1.1",
"PORT STATE SERVICE",
"22/tcp open ssh",
"80/tcp open http",
"Nmap done: 1 IP address (1 host up) scanned in 2.34 seconds",
"root@kali:~# " # Final prompt
]
current_output = []
def capture_pane_side_effect(*args, **kwargs):
# Simulate output growing over time
if len(current_output) < len(nmap_outputs):
current_output.append(nmap_outputs[len(current_output)])
return {"status": "success", "output": "\n".join(current_output)}
tmux_manager.capture_pane.side_effect = capture_pane_side_effect
trigger_system = TriggerSystem(ssh_manager, tmux_manager)
# Start monitoring for nmap completion
start_time = time.time()
result = await trigger_system.monitor_session_with_triggers(
session_id="nmap_test",
triggers=[
{"type": "regex", "pattern": "Nmap done", "name": "scan_complete"},
{"type": "timeout", "timeout_seconds": 30}
],
max_timeout=10, # Short timeout for test
poll_interval=0.1 # Fast polling for test
)
assert result["status"] == "trigger_matched"
assert result["trigger"]["trigger_name"] == "scan_complete"
assert "Nmap done" in result["output"]
print("Nmap scan scenario works with completion trigger")
@pytest.mark.asyncio
async def test_metasploit_interactive(self, mock_managers):
"""Test metasploit interactive session with prompt detection."""
ssh_manager, tmux_manager = mock_managers
# Simulate metasploit startup
msf_outputs = [
"Starting Metasploit Framework...",
"msf6 > " # Metasploit prompt
]
current_output = []
def capture_pane_side_effect(*args, **kwargs):
if len(current_output) < len(msf_outputs):
current_output.append(msf_outputs[len(current_output)])
return {"status": "success", "output": "\n".join(current_output)}
tmux_manager.capture_pane.side_effect = capture_pane_side_effect
trigger_system = TriggerSystem(ssh_manager, tmux_manager)
result = await trigger_system.monitor_session_with_triggers(
session_id="msf_test",
triggers=[
{"type": "regex", "pattern": "msf6 >", "name": "msf_ready"},
{"type": "timeout", "timeout_seconds": 10}
],
max_timeout=5,
poll_interval=0.1
)
assert result["status"] == "trigger_matched"
assert result["trigger"]["trigger_name"] == "msf_ready"
assert "msf6 >" in result["output"]
print("Metasploit interactive session works with prompt detection")
async def run_all_tests():
"""Run all trigger system tests."""
print("Running Phase 2 Trigger System Tests...\n")
# Run unit tests
test_system = TestTriggerSystem()
# Create mock managers for the tests
ssh_manager = Mock(spec=SSHManager)
tmux_manager = Mock(spec=TmuxManager)
trigger_system = TriggerSystem(ssh_manager, tmux_manager)
test_system.test_prompt_detection()
test_system.test_regex_trigger()
test_system.test_timeout_trigger()
await test_system.test_file_exists_trigger(trigger_system, tmux_manager)
test_system.test_error_detection()
test_system.test_output_buffer()
test_system.test_trigger_validation()
# Run integration tests
test_scenarios = TestRealWorldScenarios()
await test_scenarios.test_nmap_scan_scenario()
await test_scenarios.test_metasploit_interactive()
print("\nAll Phase 2 tests passed! Trigger system is ready.")
if __name__ == "__main__":
asyncio.run(run_all_tests())
