authorization_engine.py

"""
Authorization Engine with RBAC, ABAC, and Policy Evaluation.

This module provides a comprehensive authorization engine that supports:
- Role-Based Access Control (RBAC)
- Attribute-Based Access Control (ABAC)
- Policy evaluation with custom conditions
"""

import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any, Callable, Dict, List, Optional, Set, Union
from functools import wraps

logger = logging.getLogger(__name__)


class Decision(Enum):
    """Authorization decision outcomes."""
    ALLOW = "allow"
    DENY = "deny"
    INDETERMINATE = "indeterminate"


@dataclass(frozen=True)
class Subject:
    """Represents the entity requesting access."""
    id: str
    roles: Set[str] = field(default_factory=set)
    attributes: Dict[str, Any] = field(default_factory=dict)

    def has_role(self, role: str) -> bool:
        """Check if subject has a specific role."""
        return role in self.roles

    def get_attribute(self, key: str, default: Any = None) -> Any:
        """Get subject attribute with optional default."""
        return self.attributes.get(key, default)


@dataclass(frozen=True)
class Resource:
    """Represents the entity being accessed."""
    id: str
    type: str
    owner_id: Optional[str] = None
    attributes: Dict[str, Any] = field(default_factory=dict)

    def get_attribute(self, key: str, default: Any = None) -> Any:
        """Get resource attribute with optional default."""
        return self.attributes.get(key, default)


@dataclass(frozen=True)
class Action:
    """Represents the action being performed."""
    name: str
    attributes: Dict[str, Any] = field(default_factory=dict)

    def get_attribute(self, key: str, default: Any = None) -> Any:
        """Get action attribute with optional default."""
        return self.attributes.get(key, default)


@dataclass
class Policy:
    """Authorization policy with conditions."""
    id: str
    effect: Decision
    description: str = ""
    conditions: List[Callable[[Subject, Resource, Action], bool]] = field(default_factory=list)
    priority: int = 0

    def evaluate(self, subject: Subject, resource: Resource, action: Action) -> Decision:
        """
        Evaluate policy conditions.
        
        Args:
            subject: The subject requesting access
            resource: The resource being accessed
            action: The action being performed
            
        Returns:
            Decision: The policy evaluation result
        """
        try:
            for condition in self.conditions:
                if not condition(subject, resource, action):
                    return Decision.INDETERMINATE
            return self.effect
        except Exception as e:
            logger.error(f"Error evaluating policy {self.id}: {e}")
            return Decision.INDETERMINATE


class PolicyStore:
    """Manages policy storage and retrieval."""
    
    def __init__(self):
        self._policies: Dict[str, Policy] = {}
        self._role_policies: Dict[str, List[str]] = {}
    
    def add_policy(self, policy: Policy) -> None:
        """Add a policy to the store."""
        self._policies[policy.id] = policy
    
    def get_policy(self, policy_id: str) -> Optional[Policy]:
        """Retrieve a policy by ID."""
        return self._policies.get(policy_id)
    
    def get_policies(self) -> List[Policy]:
        """Get all policies."""
        return list(self._policies.values())
    
    def assign_policy_to_role(self, role: str, policy_id: str) -> None:
        """Assign a policy to a role."""
        if role not in self._role_policies:
            self._role_policies[role] = []
        self._role_policies[role].append(policy_id)
    
    def get_role_policies(self, role: str) -> List[Policy]:
        """Get policies assigned to a role."""
        policy_ids = self._role_policies.get(role, [])
        return [self._policies[pid] for pid in policy_ids if pid in self._policies]


class RBACManager:
    """Role-Based Access Control manager."""
    
    def __init__(self):
        self._role_hierarchy: Dict[str, Set[str]] = {}
        self._permissions: Dict[str, Set[str]] = {}
    
    def add_role(self, role: str, parent_roles: Optional[Set[str]] = None) -> None:
        """
        Add a role with optional parent roles.
        
        Args:
            role: The role to add
            parent_roles: Set of parent roles that inherit permissions
        """
        if parent_roles:
            if role not in self._role_hierarchy:
                self._role_hierarchy[role] = set()
            self._role_hierarchy[role].update(parent_roles)
    
    def add_permission_to_role(self, role: str, permission: str) -> None:
        """Add a permission to a role."""
        if role not in self._permissions:
            self._permissions[role] = set()
        self._permissions[role].add(permission)
    
    def get_role_permissions(self, role: str) -> Set[str]:
        """Get all permissions for a role including inherited ones."""
        permissions = set(self._permissions.get(role, set()))
        
        # Add permissions from parent roles
        for parent_role in self._role_hierarchy.get(role, set()):
            permissions.update(self.get_role_permissions(parent_role))
        
        return permissions
    
    def can_role_perform_action(self, role: str, action: str) -> bool:
        """Check if a role can perform an action."""
        return action in self.get_role_permissions(role)


class ABACCondition:
    """Predefined ABAC conditions for common use cases."""
    
    @staticmethod
    def subject_has_role(role: str) -> Callable[[Subject, Resource, Action], bool]:
        """Condition: subject has specific role."""
        def condition(subject: Subject, resource: Resource, action: Action) -> bool:
            return subject.has_role(role)
        return condition
    
    @staticmethod
    def resource_owner() -> Callable[[Subject, Resource, Action], bool]:
        """Condition: subject is resource owner."""
        def condition(subject: Subject, resource: Resource, action: Action) -> bool:
            return resource.owner_id == subject.id
        return condition
    
    @staticmethod
    def resource_attribute_equals(attr_name: str, value: Any) -> Callable[[Subject, Resource, Action], bool]:
        """Condition: resource attribute equals specific value."""
        def condition(subject: Subject, resource: Resource, action: Action) -> bool:
            return resource.get_attribute(attr_name) == value
        return condition
    
    @staticmethod
    def subject_attribute_equals(attr_name: str, value: Any) -> Callable[[Subject, Resource, Action], bool]:
        """Condition: subject attribute equals specific value."""
        def condition(subject: Subject, resource: Resource, action: Action) -> bool:
            return subject.get_attribute(attr_name) == value
        return condition
    
    @staticmethod
    def time_based(start_time: datetime, end_time: datetime) -> Callable[[Subject, Resource, Action], bool]:
        """Condition: current time is within specified range."""
        def condition(subject: Subject, resource: Resource, action: Action) -> bool:
            now = datetime.now()
            return start_time <= now <= end_time
        return condition


class AuthorizationEngine:
    """Main authorization engine combining RBAC, ABAC, and policy evaluation."""
    
    def __init__(self):
        self.policy_store = PolicyStore()
        self.rbac_manager = RBACManager()
        self._default_decision = Decision.DENY
    
    def set_default_decision(self, decision: Decision) -> None:
        """Set the default decision when no policies match."""
        self._default_decision = decision
    
    def authorize(self, subject: Subject, resource: Resource, action: Action) -> Decision:
        """
        Authorize a subject to perform an action on a resource.
        
        Args:
            subject: The subject requesting access
            resource: The resource being accessed
            action: The action being performed
            
        Returns:
            Decision: The authorization decision
        """
        logger.debug(f"Authorizing {subject.id} to {action.name} on {resource.id}")
        
        # First check RBAC
        rbac_decision = self._evaluate_rbac(subject, resource, action)
        if rbac_decision != Decision.INDETERMINATE:
            logger.debug(f"RBAC decision: {rbac_decision}")
            return rbac_decision
        
        # Then evaluate policies
        policy_decision = self._evaluate_policies(subject, resource, action)
        logger.debug(f"Policy decision: {policy_decision}")
        
        return policy_decision if policy_decision != Decision.INDETERMINATE else self._default_decision
    
    def _evaluate_rbac(self, subject: Subject, resource: Resource, action: Action) -> Decision:
        """Evaluate RBAC rules."""
        try:
            # Check if any of the subject's roles can perform the action
            for role in subject.roles:
                if self.rbac_manager.can_role_perform_action(role, action.name):
                    return Decision.ALLOW
            return Decision.INDETERMINATE
        except Exception as e:
            logger.error(f"Error in RBAC evaluation: {e}")
            return Decision.INDETERMINATE
    
    def _evaluate_policies(self, subject: Subject, resource: Resource, action: Action) -> Decision:
        """Evaluate policies in order of priority."""
        try:
            # Get all relevant policies
            policies = []
            
            # Add role-specific policies
            for role in subject.roles:
                policies.extend(self.policy_store.get_role_policies(role))
            
            # Add general policies
            policies.extend(self.policy_store.get_policies())
            
            # Sort by priority (higher priority first)
            policies.sort(key=lambda p: p.priority, reverse=True)
            
            # Evaluate policies
            for policy in policies:
                decision = policy.evaluate(subject, resource, action)
                if decision in [Decision.ALLOW, Decision.DENY]:
                    return decision
            
            return Decision.INDETERMINATE
        except Exception as e:
            logger.error(f"Error in policy evaluation: {e}")
            return Decision.INDETERMINATE


def require_authorization(engine: AuthorizationEngine, action_name: str = None):
    """
    Decorator for authorization checks.
    
    Args:
        engine: The authorization engine to use
        action_name: The action name (defaults to function name)
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            # Extract authorization context from function arguments
            # This is a simplified approach - in practice, you'd want more sophisticated context extraction
            subject = kwargs.get('subject')
            resource = kwargs.get('resource')
            
            if not subject or not resource:
                raise ValueError("Subject and resource must be provided for authorization")
            
            action = Action(action_name or func.__name__)
            decision = engine.authorize(subject, resource, action)
            
            if decision != Decision.ALLOW:
                raise PermissionError(f"Access denied: {decision.value}")
            
            return func(*args, **kwargs)
        return wrapper
    return decorator


# Example usage and configuration
def create_example_engine() -> AuthorizationEngine:
    """Create an example authorization engine with sample policies."""
    engine = AuthorizationEngine()
    
    # Setup RBAC
    engine.rbac_manager.add_role("admin")
    engine.rbac_manager.add_role("user")
    engine.rbac_manager.add_role("manager", {"user"})
    
    engine.rbac_manager.add_permission_to_role("admin", "read")
    engine.rbac_manager.add_permission_to_role("admin", "write")
    engine.rbac_manager.add_permission_to_role("admin", "delete")
    engine.rbac_manager.add_permission_to_role("user", "read")
    engine.rbac_manager.add_permission_to_role("manager", "write")
    
    # Setup policies
    # Admin can do anything
    admin_policy = Policy(
        id="admin_access",
        effect=Decision.ALLOW,
        conditions=[ABACCondition.subject_has_role("admin")],
        priority=100
    )
    
    # Users can only read published documents
    user_read_policy = Policy(
        id="user_read_published",
        effect=Decision.ALLOW,
        conditions=[
            ABACCondition.subject_has_role("user"),
            ABACCondition.resource_attribute_equals("published", True)
        ],
        priority=50
    )
    
    # Time-based access restriction
    time_restriction_policy = Policy(
        id="time_restriction",
        effect=Decision.DENY,
        conditions=[
            ABACCondition.time_based(
                datetime(2023, 1, 1, 0, 0, 0),
                datetime(2023, 12, 31, 23, 59, 59)
            )
        ],
        priority=75
    )
    
    # Add policies to store
    engine.policy_store.add_policy(admin_policy)
    engine.policy_store.add_policy(user_read_policy)
    engine.policy_store.add_policy(time_restriction_policy)
    
    # Assign policies to roles
    engine.policy_store.assign_policy_to_role("admin", "admin_access")
    engine.policy_store.assign_policy_to_role("user", "user_read_published")
    
    return engine


if __name__ == "__main__":
    # Example usage
    logging.basicConfig(level=logging.DEBUG)
    
    engine = create_example_engine()
    
    # Create test subjects
    admin = Subject(
        id="admin123",
        roles={"admin"},
        attributes={"department": "IT"}
    )
    
    user = Subject(
        id="user456",
        roles={"user"},
        attributes={"department": "HR"}
    )
    
    # Create test resources
    published_doc = Resource(
        id="doc1",
        type="document",
        owner_id="user456",
        attributes={"published": True, "sensitive": False}
    )
    
    private_doc = Resource(
        id="doc2",
        type="document",
        owner_id="user456",
        attributes={"published": False, "sensitive": True}
    )
    
    # Test authorizations
    print("Admin reading published doc:", engine.authorize(admin, published_doc, Action("read")))
    print("User reading published doc:", engine.authorize(user, published_doc, Action("read")))
    print("User reading private doc:", engine.authorize(user, private_doc, Action("read")))
    print("User writing published doc:", engine.authorize(user, published_doc, Action("write")))