breakpilot-pwa/ai-compliance-sdk/internal/rbac/policy_engine.go

package rbac

import (
	"context"
	"strings"

	"github.com/google/uuid"
)

// PolicyEngine provides advanced policy evaluation
type PolicyEngine struct {
	service *Service
	store   *Store
}

// NewPolicyEngine creates a new policy engine
func NewPolicyEngine(service *Service, store *Store) *PolicyEngine {
	return &PolicyEngine{
		service: service,
		store:   store,
	}
}

// LLMAccessRequest represents a request to access LLM functionality
type LLMAccessRequest struct {
	UserID          uuid.UUID
	TenantID        uuid.UUID
	NamespaceID     *uuid.UUID
	Model           string
	DataCategories  []string
	TokensRequested int
	Operation       string // "query", "completion", "embedding", "analysis"
}

// LLMAccessResult represents the result of an LLM access check
type LLMAccessResult struct {
	Allowed           bool
	Reason            string
	Policy            *LLMPolicy
	RequirePIIRedaction bool
	PIIRedactionLevel PIIRedactionLevel
	MaxTokens         int
	BlockedCategories []string
}

// EvaluateLLMAccess evaluates whether an LLM request should be allowed
func (pe *PolicyEngine) EvaluateLLMAccess(ctx context.Context, req *LLMAccessRequest) (*LLMAccessResult, error) {
	result := &LLMAccessResult{
		Allowed:           false,
		RequirePIIRedaction: true,
		PIIRedactionLevel: PIIRedactionStrict,
		MaxTokens:         4000,
		BlockedCategories: []string{},
	}

	// 1. Check base permission for LLM access
	hasPermission, err := pe.service.HasAnyPermission(ctx, req.UserID, req.TenantID, req.NamespaceID, []string{
		PermissionLLMAll,
		PermissionLLMQuery,
		PermissionLLMOwnQuery,
	})
	if err != nil {
		return result, err
	}
	if !hasPermission {
		result.Reason = "no LLM permission"
		return result, nil
	}

	// 2. Get effective policy
	policy, err := pe.store.GetEffectiveLLMPolicy(ctx, req.TenantID, req.NamespaceID)
	if err != nil {
		return result, err
	}

	result.Policy = policy

	// No policy = use defaults and allow
	if policy == nil {
		result.Allowed = true
		result.Reason = "no policy restrictions"
		return result, nil
	}

	// 3. Check model restrictions
	if len(policy.AllowedModels) > 0 {
		modelAllowed := false
		for _, allowed := range policy.AllowedModels {
			if allowed == req.Model || strings.HasPrefix(req.Model, allowed+":") || strings.HasPrefix(req.Model, allowed+"-") {
				modelAllowed = true
				break
			}
		}
		if !modelAllowed {
			result.Reason = "model not allowed: " + req.Model
			return result, nil
		}
	}

	// 4. Check data categories
	for _, category := range req.DataCategories {
		// Check if blocked
		for _, blocked := range policy.BlockedDataCategories {
			if blocked == category {
				result.BlockedCategories = append(result.BlockedCategories, category)
			}
		}

		// Check if allowed (if allowlist is defined)
		if len(policy.AllowedDataCategories) > 0 {
			allowed := false
			for _, a := range policy.AllowedDataCategories {
				if a == category {
					allowed = true
					break
				}
			}
			if !allowed {
				result.BlockedCategories = append(result.BlockedCategories, category)
			}
		}
	}

	if len(result.BlockedCategories) > 0 {
		result.Reason = "blocked data categories: " + strings.Join(result.BlockedCategories, ", ")
		return result, nil
	}

	// 5. Check token limits
	if req.TokensRequested > policy.MaxTokensPerRequest {
		result.Reason = "tokens requested exceeds limit"
		return result, nil
	}
	result.MaxTokens = policy.MaxTokensPerRequest

	// 6. Set PII redaction requirements
	result.RequirePIIRedaction = policy.RequirePIIRedaction
	result.PIIRedactionLevel = policy.PIIRedactionLevel

	// All checks passed
	result.Allowed = true
	result.Reason = "policy check passed"
	return result, nil
}

// NamespaceAccessRequest represents a request to access a namespace
type NamespaceAccessRequest struct {
	UserID      uuid.UUID
	TenantID    uuid.UUID
	NamespaceID uuid.UUID
	Operation   string // "read", "write", "admin"
}

// NamespaceAccessResult represents the result of a namespace access check
type NamespaceAccessResult struct {
	Allowed            bool
	Reason             string
	DataClassification DataClassification
	IsolationLevel     IsolationLevel
	Permissions        []string
}

// EvaluateNamespaceAccess evaluates whether a namespace operation should be allowed
func (pe *PolicyEngine) EvaluateNamespaceAccess(ctx context.Context, req *NamespaceAccessRequest) (*NamespaceAccessResult, error) {
	result := &NamespaceAccessResult{
		Allowed:     false,
		Permissions: []string{},
	}

	// Get namespace details
	ns, err := pe.store.GetNamespace(ctx, req.NamespaceID)
	if err != nil {
		result.Reason = "namespace not found"
		return result, ErrNamespaceNotFound
	}

	result.DataClassification = ns.DataClassification
	result.IsolationLevel = ns.IsolationLevel

	// Check if user has any roles in this namespace
	userRoles, err := pe.store.GetUserRolesForNamespace(ctx, req.UserID, req.TenantID, &req.NamespaceID)
	if err != nil {
		return result, err
	}

	if len(userRoles) == 0 {
		// Check for tenant-wide roles
		tenantRoles, err := pe.store.GetUserRolesForNamespace(ctx, req.UserID, req.TenantID, nil)
		if err != nil {
			return result, err
		}

		if len(tenantRoles) == 0 {
			result.Reason = "no access to namespace"
			return result, nil
		}

		userRoles = tenantRoles
	}

	// Collect permissions
	permSet := make(map[string]bool)
	for _, ur := range userRoles {
		for _, perm := range ur.RolePermissions {
			permSet[perm] = true
		}
	}

	for perm := range permSet {
		result.Permissions = append(result.Permissions, perm)
	}

	// Check operation-specific permission
	var requiredPermission string
	switch req.Operation {
	case "read":
		requiredPermission = "namespace:read"
	case "write":
		requiredPermission = "namespace:write"
	case "admin":
		requiredPermission = "namespace:own:admin"
	default:
		requiredPermission = "namespace:read"
	}

	hasPermission := pe.service.checkPermission(result.Permissions, requiredPermission)
	if !hasPermission {
		// Check for broader permissions
		hasPermission = pe.service.checkPermission(result.Permissions, PermissionComplianceAll) ||
			pe.service.checkPermission(result.Permissions, "namespace:*")
	}

	if !hasPermission {
		result.Reason = "insufficient permissions for operation: " + req.Operation
		return result, nil
	}

	result.Allowed = true
	result.Reason = "access granted"
	return result, nil
}

// DataAccessRequest represents a request to access data
type DataAccessRequest struct {
	UserID         uuid.UUID
	TenantID       uuid.UUID
	NamespaceID    *uuid.UUID
	ResourceType   string   // "control", "evidence", "audit", "policy"
	ResourceID     *uuid.UUID
	Operation      string   // "read", "create", "update", "delete"
	DataCategories []string
}

// DataAccessResult represents the result of a data access check
type DataAccessResult struct {
	Allowed           bool
	Reason            string
	RequireAuditLog   bool
	AllowedCategories []string
	DeniedCategories  []string
}

// EvaluateDataAccess evaluates whether a data operation should be allowed
func (pe *PolicyEngine) EvaluateDataAccess(ctx context.Context, req *DataAccessRequest) (*DataAccessResult, error) {
	result := &DataAccessResult{
		Allowed:           false,
		RequireAuditLog:   true,
		AllowedCategories: []string{},
		DeniedCategories:  []string{},
	}

	// Build required permission based on resource type and operation
	requiredPermission := req.ResourceType + ":" + req.Operation

	// Check permission
	hasPermission, err := pe.service.HasAnyPermission(ctx, req.UserID, req.TenantID, req.NamespaceID, []string{
		requiredPermission,
		req.ResourceType + ":*",
		PermissionComplianceAll,
	})
	if err != nil {
		return result, err
	}

	if !hasPermission {
		result.Reason = "missing permission: " + requiredPermission
		return result, nil
	}

	// Check namespace access if namespace is specified
	if req.NamespaceID != nil {
		nsResult, err := pe.EvaluateNamespaceAccess(ctx, &NamespaceAccessRequest{
			UserID:      req.UserID,
			TenantID:    req.TenantID,
			NamespaceID: *req.NamespaceID,
			Operation:   req.Operation,
		})
		if err != nil {
			return result, err
		}

		if !nsResult.Allowed {
			result.Reason = "namespace access denied: " + nsResult.Reason
			return result, nil
		}
	}

	// Check data categories
	if len(req.DataCategories) > 0 {
		policy, _ := pe.store.GetEffectiveLLMPolicy(ctx, req.TenantID, req.NamespaceID)
		if policy != nil {
			for _, category := range req.DataCategories {
				blocked := false

				// Check blocked list
				for _, b := range policy.BlockedDataCategories {
					if b == category {
						blocked = true
						break
					}
				}

				// Check allowed list if specified
				if !blocked && len(policy.AllowedDataCategories) > 0 {
					allowed := false
					for _, a := range policy.AllowedDataCategories {
						if a == category {
							allowed = true
							break
						}
					}
					if !allowed {
						blocked = true
					}
				}

				if blocked {
					result.DeniedCategories = append(result.DeniedCategories, category)
				} else {
					result.AllowedCategories = append(result.AllowedCategories, category)
				}
			}

			if len(result.DeniedCategories) > 0 {
				result.Reason = "blocked data categories: " + strings.Join(result.DeniedCategories, ", ")
				return result, nil
			}
		}
	}

	result.Allowed = true
	result.Reason = "access granted"
	return result, nil
}

// GetUserContext returns full context for a user including all permissions and policies
func (pe *PolicyEngine) GetUserContext(ctx context.Context, userID, tenantID uuid.UUID) (*UserContext, error) {
	perms, err := pe.service.GetEffectivePermissions(ctx, userID, tenantID, nil)
	if err != nil {
		return nil, err
	}

	tenant, err := pe.store.GetTenant(ctx, tenantID)
	if err != nil {
		return nil, err
	}

	namespaces, err := pe.service.GetUserAccessibleNamespaces(ctx, userID, tenantID)
	if err != nil {
		return nil, err
	}

	return &UserContext{
		UserID:      userID,
		TenantID:    tenantID,
		TenantName:  tenant.Name,
		TenantSlug:  tenant.Slug,
		Roles:       perms.Roles,
		Permissions: perms.Permissions,
		Namespaces:  namespaces,
		LLMPolicy:   perms.LLMPolicy,
	}, nil
}

// UserContext represents complete context for a user
type UserContext struct {
	UserID      uuid.UUID    `json:"user_id"`
	TenantID    uuid.UUID    `json:"tenant_id"`
	TenantName  string       `json:"tenant_name"`
	TenantSlug  string       `json:"tenant_slug"`
	Roles       []string     `json:"roles"`
	Permissions []string     `json:"permissions"`
	Namespaces  []*Namespace `json:"namespaces"`
	LLMPolicy   *LLMPolicy   `json:"llm_policy,omitempty"`
}