rpa_vision_v3/test_upload_with_hardcoded_token.py

#!/usr/bin/env python3
"""
Real Functionality Test: Authentication and Upload System

Tests the complete authentication and upload flow using real components:
- Real FastAPI server with security middleware
- Real token validation system
- Real file upload and processing
- Real RawSession creation and validation
- Real storage system integration

This test validates actual business logic without mocks or simulations.
"""

import json
import os
import sys
import tempfile
import time
import zipfile
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import pytest
import requests

# Add project root to path for imports
sys.path.insert(0, str(Path(__file__).parent))

from core.models.raw_session import RawSession
from core.security.api_tokens import get_token_manager, TokenRole


class RealUploadTest:
    """
    Real functionality test for the upload authentication system.
    
    Tests the complete flow from token validation to file processing
    using actual system components.
    """
    
    def __init__(self):
        self.server_url = "http://127.0.0.1:8000"
        self.upload_endpoint = f"{self.server_url}/api/traces/upload"
        self.status_endpoint = f"{self.server_url}/api/traces/status"
        self.debug_auth_endpoint = f"{self.server_url}/api/traces/debug-auth"
        
        # Load tokens from environment and system
        self.tokens = self._load_real_tokens()
        
        # Test data directories
        self.test_data_dir = Path("data/test_uploads")
        self.test_data_dir.mkdir(parents=True, exist_ok=True)
    
    def _load_real_tokens(self) -> Dict[str, str]:
        """Load real tokens from environment and token manager."""
        tokens = {}
        
        # Load from environment (.env.local)
        env_admin = os.getenv("RPA_TOKEN_ADMIN")
        env_readonly = os.getenv("RPA_TOKEN_READONLY")
        
        if env_admin:
            tokens["env_admin"] = env_admin
        if env_readonly:
            tokens["env_readonly"] = env_readonly
        
        # Get tokens from the actual token manager (includes hardcoded production tokens)
        token_manager = get_token_manager()
        
        # Extract some admin tokens for testing
        admin_tokens = list(token_manager.admin_tokens)[:3]  # Take first 3
        readonly_tokens = list(token_manager.read_only_tokens)[:3]  # Take first 3
        
        for i, token in enumerate(admin_tokens):
            tokens[f"admin_{i}"] = token
        
        for i, token in enumerate(readonly_tokens):
            tokens[f"readonly_{i}"] = token
        
        return tokens
    
    def _create_real_raw_session(self, session_id: str) -> Tuple[RawSession, Path]:
        """
        Create a real RawSession with valid structure and data.
        
        Returns:
            Tuple of (RawSession object, path to created ZIP file)
        """
        # Create realistic session data
        session = RawSession(
            schema_version="rawsession_v1",
            session_id=session_id,
            agent_version="0.1.0",
            environment={
                "platform": "linux",
                "hostname": "test-machine",
                "screen": {
                    "primary_resolution": [1920, 1080],
                    "display_scale": 1.0
                }
            },
            user={
                "id": "test_user",
                "label": "Test User"
            },
            context={
                "customer": "Test Company",
                "training_label": "Authentication Test",
                "notes": "Real functionality test session"
            },
            started_at=datetime.now(),
            ended_at=datetime.now(),
            events=[
                {
                    "t": 0.5,
                    "type": "mouse_click",
                    "button": "left",
                    "pos": [100, 200],
                    "window": {
                        "title": "Test Application",
                        "app_name": "test_app"
                    },
                    "screenshot_id": "shot_0001"
                },
                {
                    "t": 1.2,
                    "type": "key_combo",
                    "keys": ["CTRL", "C"],
                    "window": {
                        "title": "Test Application",
                        "app_name": "test_app"
                    },
                    "screenshot_id": "shot_0002"
                }
            ],
            screenshots=[
                {
                    "screenshot_id": "shot_0001",
                    "relative_path": "shots/shot_0001.png",
                    "captured_at": datetime.now().isoformat()
                },
                {
                    "screenshot_id": "shot_0002",
                    "relative_path": "shots/shot_0002.png",
                    "captured_at": datetime.now().isoformat()
                }
            ]
        )
        
        # Create session directory structure
        session_dir = self.test_data_dir / session_id / session_id
        session_dir.mkdir(parents=True, exist_ok=True)
        
        # Create shots directory
        shots_dir = session_dir / "shots"
        shots_dir.mkdir(exist_ok=True)
        
        # Create dummy screenshot files
        for screenshot in session.screenshots:
            shot_path = session_dir / screenshot["relative_path"]
            shot_path.write_bytes(b"fake_png_data_for_testing")
        
        # Save session JSON
        json_path = session_dir / f"{session_id}.json"
        session.save_to_file(json_path)
        
        # Create ZIP file
        zip_path = self.test_data_dir / f"{session_id}.zip"
        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:
            # Add all files in session directory
            for file_path in session_dir.rglob("*"):
                if file_path.is_file():
                    arcname = file_path.relative_to(self.test_data_dir)
                    zf.write(file_path, arcname)
        
        return session, zip_path
    
    def _check_server_availability(self) -> bool:
        """Check if the server is running and accessible."""
        try:
            response = requests.get(self.status_endpoint, timeout=5)
            return response.status_code == 200
        except requests.RequestException:
            return False
    
    def _test_token_validation(self, token: str, token_name: str) -> Dict:
        """Test token validation using the real token manager."""
        print(f"\n--- Testing Token Validation: {token_name} ---")
        
        token_manager = get_token_manager()
        
        try:
            token_info = token_manager.validate_token(token)
            result = {
                "valid": True,
                "role": token_info.role.value,
                "user_id": token_info.user_id,
                "metadata": token_info.metadata
            }
            print(f"✅ Token valid: {result}")
            return result
        except Exception as e:
            result = {
                "valid": False,
                "error": str(e)
            }
            print(f"❌ Token invalid: {result}")
            return result
    
    def _test_upload_with_token(self, token: str, token_name: str, 
                               session_id: str, zip_path: Path) -> Dict:
        """Test upload with a specific token using real HTTP request."""
        print(f"\n--- Testing Upload: {token_name} ---")
        
        test_results = []
        
        # Test 1: X-Admin-Token header
        print(f"1. Testing X-Admin-Token header...")
        with open(zip_path, 'rb') as f:
            files = {'file': (f'{session_id}.zip', f, 'application/zip')}
            data = {'session_id': session_id}
            headers = {'X-Admin-Token': token}
            
            try:
                response = requests.post(
                    self.upload_endpoint, 
                    files=files, 
                    data=data, 
                    headers=headers,
                    timeout=30
                )
                result = {
                    "method": "X-Admin-Token",
                    "status_code": response.status_code,
                    "success": response.status_code == 200,
                    "response": response.text[:500] if response.text else None
                }
                test_results.append(result)
                print(f"   Status: {response.status_code}")
                print(f"   Success: {result['success']}")
                if not result['success']:
                    print(f"   Error: {response.text}")
            except requests.RequestException as e:
                result = {
                    "method": "X-Admin-Token",
                    "status_code": None,
                    "success": False,
                    "error": str(e)
                }
                test_results.append(result)
                print(f"   Request failed: {e}")
        
        # Test 2: Authorization Bearer header
        print(f"2. Testing Authorization Bearer header...")
        with open(zip_path, 'rb') as f:
            files = {'file': (f'{session_id}.zip', f, 'application/zip')}
            data = {'session_id': session_id}
            headers = {'Authorization': f'Bearer {token}'}
            
            try:
                response = requests.post(
                    self.upload_endpoint, 
                    files=files, 
                    data=data, 
                    headers=headers,
                    timeout=30
                )
                result = {
                    "method": "Authorization Bearer",
                    "status_code": response.status_code,
                    "success": response.status_code == 200,
                    "response": response.text[:500] if response.text else None
                }
                test_results.append(result)
                print(f"   Status: {response.status_code}")
                print(f"   Success: {result['success']}")
                if not result['success']:
                    print(f"   Error: {response.text}")
            except requests.RequestException as e:
                result = {
                    "method": "Authorization Bearer",
                    "status_code": None,
                    "success": False,
                    "error": str(e)
                }
                test_results.append(result)
                print(f"   Request failed: {e}")
        
        return {
            "token_name": token_name,
            "tests": test_results,
            "any_success": any(r["success"] for r in test_results)
        }
    
    def _test_unauthorized_access(self) -> Dict:
        """Test that requests without tokens are properly rejected."""
        print(f"\n--- Testing Unauthorized Access ---")
        
        session_id = "unauthorized_test"
        _, zip_path = self._create_real_raw_session(session_id)
        
        try:
            with open(zip_path, 'rb') as f:
                files = {'file': (f'{session_id}.zip', f, 'application/zip')}
                data = {'session_id': session_id}
                # No authentication headers
                
                response = requests.post(
                    self.upload_endpoint, 
                    files=files, 
                    data=data,
                    timeout=30
                )
                
                result = {
                    "status_code": response.status_code,
                    "rejected": response.status_code in [401, 403],
                    "response": response.text[:200] if response.text else None
                }
                
                print(f"Status: {response.status_code}")
                print(f"Properly rejected: {result['rejected']}")
                if not result['rejected']:
                    print(f"⚠️  Security issue: Unauthorized request was not rejected!")
                
                return result
                
        except requests.RequestException as e:
            return {
                "status_code": None,
                "rejected": True,  # Network error counts as rejection
                "error": str(e)
            }
    
    def _verify_server_debug_endpoints(self) -> Dict:
        """Test server debug endpoints to verify token configuration."""
        print(f"\n--- Testing Server Debug Endpoints ---")
        
        try:
            response = requests.get(self.debug_auth_endpoint, timeout=10)
            if response.status_code == 200:
                debug_info = response.json()
                print(f"✅ Debug endpoint accessible")
                print(f"Admin tokens configured: {debug_info.get('token_counts', {}).get('admin_tokens', 0)}")
                print(f"Read-only tokens configured: {debug_info.get('token_counts', {}).get('read_only_tokens', 0)}")
                return {
                    "accessible": True,
                    "debug_info": debug_info
                }
            else:
                print(f"❌ Debug endpoint returned {response.status_code}")
                return {
                    "accessible": False,
                    "status_code": response.status_code
                }
        except requests.RequestException as e:
            print(f"❌ Debug endpoint not accessible: {e}")
            return {
                "accessible": False,
                "error": str(e)
            }
    
    def run_complete_test(self) -> Dict:
        """
        Run the complete real functionality test suite.
        
        Returns:
            Complete test results
        """
        print("🚀 Starting Real Functionality Test: Authentication & Upload System")
        print("=" * 70)
        
        # Check server availability
        if not self._check_server_availability():
            return {
                "success": False,
                "error": "Server not available at " + self.server_url,
                "suggestion": "Start the server with: python server/api_upload.py"
            }
        
        print(f"✅ Server is running at {self.server_url}")
        
        results = {
            "server_available": True,
            "tokens_tested": {},
            "unauthorized_test": {},
            "debug_endpoints": {},
            "summary": {}
        }
        
        # Test debug endpoints
        results["debug_endpoints"] = self._verify_server_debug_endpoints()
        
        # Test each token
        successful_tokens = []
        failed_tokens = []
        
        for token_name, token in self.tokens.items():
            # Validate token using real token manager
            validation_result = self._test_token_validation(token, token_name)
            
            # Test upload with token
            session_id = f"test_session_{token_name}_{int(time.time())}"
            session, zip_path = self._create_real_raw_session(session_id)
            
            upload_result = self._test_upload_with_token(token, token_name, session_id, zip_path)
            
            results["tokens_tested"][token_name] = {
                "validation": validation_result,
                "upload": upload_result
            }
            
            if upload_result["any_success"]:
                successful_tokens.append(token_name)
            else:
                failed_tokens.append(token_name)
        
        # Test unauthorized access
        results["unauthorized_test"] = self._test_unauthorized_access()
        
        # Generate summary
        results["summary"] = {
            "total_tokens_tested": len(self.tokens),
            "successful_tokens": len(successful_tokens),
            "failed_tokens": len(failed_tokens),
            "successful_token_names": successful_tokens,
            "failed_token_names": failed_tokens,
            "security_working": results["unauthorized_test"].get("rejected", False),
            "overall_success": len(successful_tokens) > 0 and results["unauthorized_test"].get("rejected", False)
        }
        
        return results
    
    def cleanup(self):
        """Clean up test files."""
        import shutil
        if self.test_data_dir.exists():
            shutil.rmtree(self.test_data_dir)


def test_real_authentication_upload():
    """
    Pytest-compatible test function for real authentication and upload functionality.
    
    This test validates:
    1. Real token validation using the actual TokenManager
    2. Real HTTP authentication with FastAPI security middleware
    3. Real file upload and processing pipeline
    4. Real RawSession creation and validation
    5. Real security enforcement (unauthorized access rejection)
    """
    test_runner = RealUploadTest()
    
    try:
        results = test_runner.run_complete_test()
        
        # Print detailed results
        print("\n" + "=" * 70)
        print("📊 TEST RESULTS SUMMARY")
        print("=" * 70)
        
        summary = results.get("summary", {})
        print(f"Total tokens tested: {summary.get('total_tokens_tested', 0)}")
        print(f"Successful tokens: {summary.get('successful_tokens', 0)}")
        print(f"Failed tokens: {summary.get('failed_tokens', 0)}")
        print(f"Security working: {summary.get('security_working', False)}")
        print(f"Overall success: {summary.get('overall_success', False)}")
        
        if summary.get('successful_token_names'):
            print(f"✅ Working tokens: {', '.join(summary['successful_token_names'])}")
        
        if summary.get('failed_token_names'):
            print(f"❌ Failed tokens: {', '.join(summary['failed_token_names'])}")
        
        # Assert test conditions
        assert results.get("server_available", False), "Server must be available"
        assert summary.get("successful_tokens", 0) > 0, "At least one token must work"
        assert summary.get("security_working", False), "Security must reject unauthorized requests"
        assert summary.get("overall_success", False), "Overall test must pass"
        
        print("\n🎉 All real functionality tests passed!")
        
    finally:
        test_runner.cleanup()


def main():
    """Main function for running the test standalone."""
    test_runner = RealUploadTest()
    
    try:
        results = test_runner.run_complete_test()
        
        # Print results
        print("\n" + "=" * 70)
        print("📊 FINAL RESULTS")
        print("=" * 70)
        print(json.dumps(results["summary"], indent=2))
        
        if results["summary"]["overall_success"]:
            print("\n🎉 All tests passed! Authentication and upload system working correctly.")
            return 0
        else:
            print("\n❌ Some tests failed. Check the detailed output above.")
            return 1
            
    except Exception as e:
        print(f"\n💥 Test execution failed: {e}")
        return 1
    finally:
        test_runner.cleanup()


if __name__ == "__main__":
    exit(main())