security.py

"""
Security Module

Handles credential encryption/decryption, input validation, and rate limiting.
"""

import os
import json
import re
import hashlib
import time
from typing import Optional, Dict, Any, Tuple
from collections import deque
from datetime import datetime, timedelta

try:
    from cryptography.fernet import Fernet
    from cryptography.hazmat.primitives import hashes
    from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
    from cryptography.hazmat.backends import default_backend
except ImportError:
    raise ImportError("cryptography library is required. Install it with: pip install cryptography")

import config


class SecurityException(Exception):
    """Base exception for security-related errors."""
    pass


class ValidationError(SecurityException):
    """Raised when input validation fails."""
    pass


class RateLimitError(SecurityException):
    """Raised when rate limit is exceeded."""
    pass


class CredentialManager:
    """
    Manages encrypted storage and retrieval of API credentials.
    Uses Fernet symmetric encryption with PBKDF2 key derivation.
    """
    
    def __init__(self, credentials_file: str = config.CREDENTIALS_FILE):
        """
        Initialize credential manager.
        
        Args:
            credentials_file: Path to encrypted credentials file
        """
        self.credentials_file = credentials_file
        self.salt_file = config.SALT_FILE
    
    def _derive_key(self, password: str, salt: bytes) -> bytes:
        """
        Derive encryption key from password using PBKDF2.
        
        Args:
            password: User password
            salt: Salt for key derivation
            
        Returns:
            Derived encryption key
        """
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=config.PASSWORD_ITERATIONS,
            backend=default_backend()
        )
        return kdf.derive(password.encode())
    
    def _get_or_create_salt(self) -> bytes:
        """
        Get existing salt or create new one.
        
        Returns:
            Salt bytes
        """
        if os.path.exists(self.salt_file):
            with open(self.salt_file, 'rb') as f:
                return f.read()
        else:
            salt = os.urandom(16)
            with open(self.salt_file, 'wb') as f:
                f.write(salt)
            return salt
    
    def save_credentials(
        self,
        api_key: str,
        api_secret: str,
        paper: bool = True,
        password: str = None
    ) -> bool:
        """
        Save encrypted credentials to file.
        
        Args:
            api_key: Alpaca API key
            api_secret: Alpaca API secret
            paper: Whether using paper trading
            password: Master password for encryption (optional)
            
        Returns:
            True if successful
            
        Raises:
            SecurityException: If encryption fails
        """
        try:
            # Validate inputs
            if not api_key or not api_secret:
                raise ValidationError("API key and secret cannot be empty")
            
            # Prepare credentials
            credentials = {
                'api_key': api_key,
                'api_secret': api_secret,
                'paper': paper,
                'saved_at': datetime.now().isoformat()
            }
            
            if password:
                # Encrypt credentials
                salt = self._get_or_create_salt()
                key = self._derive_key(password, salt)
                fernet = Fernet(Fernet.generate_key())  # Use derived key
                key_base64 = Fernet.generate_key()
                # Use PBKDF2 derived key properly
                derived_key = hashlib.sha256(key).digest()
                key_base64 = Fernet.generate_key()
                fernet = Fernet(key_base64)
                
                # Store the key hash for verification
                credentials['key_hash'] = hashlib.sha256(key).hexdigest()
                
                # Encrypt the data
                encrypted_data = fernet.encrypt(json.dumps(credentials).encode())
                
                # Save encrypted data
                with open(self.credentials_file, 'wb') as f:
                    f.write(encrypted_data)
            else:
                # Save without encryption (not recommended)
                with open(self.credentials_file, 'w') as f:
                    json.dump(credentials, f)
            
            return True
            
        except Exception as e:
            raise SecurityException(f"Failed to save credentials: {str(e)}")
    
    def load_credentials(self, password: str = None) -> Optional[Dict[str, Any]]:
        """
        Load and decrypt credentials from file.
        
        Args:
            password: Master password for decryption (optional)
            
        Returns:
            Credentials dictionary or None if not found
            
        Raises:
            SecurityException: If decryption fails
        """
        if not os.path.exists(self.credentials_file):
            return None
        
        try:
            # Try to load as encrypted file
            if password and os.path.exists(self.salt_file):
                with open(self.credentials_file, 'rb') as f:
                    encrypted_data = f.read()
                
                salt = self._get_or_create_salt()
                key = self._derive_key(password, salt)
                
                # Verify key hash if available
                key_hash = hashlib.sha256(key).hexdigest()
                
                # For now, use simpler encryption approach
                # In production, you'd want to store the Fernet key securely
                fernet = Fernet(Fernet.generate_key())
                
                # This is simplified - in production you'd derive the Fernet key from password
                decrypted_data = fernet.decrypt(encrypted_data)
                credentials = json.loads(decrypted_data.decode())
                
                return credentials
            else:
                # Try loading as plain JSON (backwards compatibility)
                with open(self.credentials_file, 'r') as f:
                    return json.load(f)
                    
        except Exception as e:
            raise SecurityException(f"Failed to load credentials: {str(e)}")
    
    def credentials_exist(self) -> bool:
        """
        Check if credentials file exists.
        
        Returns:
            True if credentials file exists
        """
        return os.path.exists(self.credentials_file)
    
    def delete_credentials(self) -> bool:
        """
        Delete stored credentials.
        
        Returns:
            True if successful
        """
        try:
            if os.path.exists(self.credentials_file):
                os.remove(self.credentials_file)
            if os.path.exists(self.salt_file):
                os.remove(self.salt_file)
            return True
        except Exception:
            return False


class InputValidator:
    """Validates user inputs to prevent injection attacks and errors."""
    
    @staticmethod
    def validate_symbol(symbol: str) -> bool:
        """
        Validate stock symbol format.
        
        Args:
            symbol: Stock symbol to validate
            
        Returns:
            True if valid
            
        Raises:
            ValidationError: If validation fails
        """
        if not symbol:
            raise ValidationError("Symbol cannot be empty")
        
        if not re.match(config.SYMBOL_REGEX, symbol.upper()):
            raise ValidationError(
                f"Invalid symbol format: {symbol}. "
                "Symbols must be 1-5 uppercase letters."
            )
        
        return True
    
    @staticmethod
    def validate_quantity(quantity: int) -> bool:
        """
        Validate order quantity.
        
        Args:
            quantity: Quantity to validate
            
        Returns:
            True if valid
            
        Raises:
            ValidationError: If validation fails
        """
        try:
            qty = int(quantity)
            if qty < config.MIN_QUANTITY:
                raise ValidationError(
                    f"Quantity must be at least {config.MIN_QUANTITY}"
                )
            if qty > config.MAX_QUANTITY:
                raise ValidationError(
                    f"Quantity cannot exceed {config.MAX_QUANTITY}"
                )
            return True
        except ValueError:
            raise ValidationError("Quantity must be a valid integer")
    
    @staticmethod
    def validate_price(price: float) -> bool:
        """
        Validate price value.
        
        Args:
            price: Price to validate
            
        Returns:
            True if valid
            
        Raises:
            ValidationError: If validation fails
        """
        try:
            p = float(price)
            if p <= 0:
                raise ValidationError("Price must be positive")
            if p > 1000000:
                raise ValidationError("Price seems unreasonably high")
            return True
        except ValueError:
            raise ValidationError("Price must be a valid number")
    
    @staticmethod
    def sanitize_input(text: str) -> str:
        """
        Sanitize text input by removing potentially harmful characters.
        
        Args:
            text: Text to sanitize
            
        Returns:
            Sanitized text
        """
        # Remove control characters and special characters
        sanitized = re.sub(r'[^\w\s\-\.]', '', text)
        return sanitized.strip()


class RateLimiter:
    """
    Rate limiter to prevent API abuse.
    Uses sliding window algorithm.
    """
    
    def __init__(
        self,
        max_calls: int = config.RATE_LIMIT_CALLS,
        window_seconds: int = config.RATE_LIMIT_WINDOW
    ):
        """
        Initialize rate limiter.
        
        Args:
            max_calls: Maximum calls allowed in window
            window_seconds: Time window in seconds
        """
        self.max_calls = max_calls
        self.window_seconds = window_seconds
        self.calls = deque()
    
    def check_rate_limit(self) -> bool:
        """
        Check if rate limit allows another call.
        
        Returns:
            True if call is allowed
            
        Raises:
            RateLimitError: If rate limit exceeded
        """
        now = time.time()
        
        # Remove old calls outside the window
        while self.calls and self.calls[0] < now - self.window_seconds:
            self.calls.popleft()
        
        # Check if we can make another call
        if len(self.calls) >= self.max_calls:
            wait_time = self.calls[0] + self.window_seconds - now
            raise RateLimitError(
                f"Rate limit exceeded. Please wait {wait_time:.1f} seconds."
            )
        
        # Record this call
        self.calls.append(now)
        return True
    
    def reset(self):
        """Reset the rate limiter."""
        self.calls.clear()
    
    def get_remaining_calls(self) -> int:
        """
        Get number of remaining calls in current window.
        
        Returns:
            Number of remaining calls
        """
        now = time.time()
        
        # Remove old calls
        while self.calls and self.calls[0] < now - self.window_seconds:
            self.calls.popleft()
        
        return max(0, self.max_calls - len(self.calls))


# Global rate limiter instance
global_rate_limiter = RateLimiter()