EonaCat.SecureToken/EonaCat.SecureToken/Extensions/TokenServiceExtensions.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using EonaCat.SecureToken.Core;
using EonaCat.SecureToken.Security;
using EonaCat.SecureToken.Tokens;
using EonaCat.SecureToken.Validation;

namespace EonaCat.SecureToken.Extensions
{
    // This file is part of the EonaCat project(s) which is released under the Apache License.
    // See the LICENSE file or go to https://EonaCat.com/License for full license details.

    /// <summary>
    /// Extension methods for advanced token operations including batch processing,
    /// security checks, and validation with additional features.
    /// </summary>
    public static class TokenServiceExtensions
    {
        /// <summary>
        /// Validates multiple tokens concurrently.
        /// </summary>
        /// <param name="service">The token service</param>
        /// <param name="tokens">Tokens to validate</param>
        /// <param name="options">Validation options</param>
        /// <param name="ct">Cancellation token</param>
        /// <returns>Collection of validation results in the same order as input tokens</returns>
        public static async Task<IReadOnlyList<TokenResult>> ValidateManyAsync(
            this ITokenService service,
            IEnumerable<string> tokens,
            TokenValidationOptions options,
            CancellationToken ct = default)
        {
            if (service is null)
            {
                throw new ArgumentNullException(nameof(service));
            }

            if (tokens is null)
            {
                throw new ArgumentNullException(nameof(tokens));
            }

            if (options is null)
            {
                throw new ArgumentNullException(nameof(options));
            }

            var tokenList = tokens.ToList();
            var results = new List<TokenResult>(tokenList.Count);

            foreach (var token in tokenList)
            {
                results.Add(await service.ValidateAsync(token, options, ct).ConfigureAwait(false));
            }

            return results.AsReadOnly();
        }

        /// <summary>
        /// Validates a token and checks the blacklist.
        /// </summary>
        /// <param name="service">The token service</param>
        /// <param name="token">Token to validate</param>
        /// <param name="options">Validation options</param>
        /// <param name="blacklist">The blacklist to check</param>
        /// <param name="ct">Cancellation token</param>
        /// <returns>Validation result (or Revoked if token is blacklisted)</returns>
        public static async Task<TokenResult> ValidateWithBlacklistAsync(
            this ITokenService service,
            string token,
            TokenValidationOptions options,
            ITokenBlacklist blacklist,
            CancellationToken ct = default)
        {
            if (service is null)
            {
                throw new ArgumentNullException(nameof(service));
            }

            if (string.IsNullOrWhiteSpace(token))
            {
                throw new ArgumentException("Token cannot be null or empty.", nameof(token));
            }

            if (blacklist is null)
            {
                throw new ArgumentNullException(nameof(blacklist));
            }

            // First validate the token normally
            var result = await service.ValidateAsync(token, options, ct).ConfigureAwait(false);

            // If validation succeeded, check blacklist
            if (result is TokenResult.Success success)
            {
                var isBlacklisted = await blacklist.IsBlacklistedAsync(success.Claims.TokenId).ConfigureAwait(false);
                if (isBlacklisted)
                {
                    return new TokenResult.Revoked(success.Claims.TokenId);
                }
            }

            return result;
        }

        /// <summary>
        /// Validates a token with rate limiting on the subject.
        /// </summary>
        /// <param name="service">The token service</param>
        /// <param name="token">Token to validate</param>
        /// <param name="options">Validation options</param>
        /// <param name="rateLimiter">The rate limiter</param>
        /// <param name="subject">Subject to rate limit (e.g., user ID, IP address)</param>
        /// <param name="ct">Cancellation token</param>
        /// <returns>Validation result (or failure if rate limited)</returns>
        public static async Task<TokenResult> ValidateWithRateLimitingAsync(
            this ITokenService service,
            string token,
            TokenValidationOptions options,
            ITokenRateLimiter rateLimiter,
            string subject,
            CancellationToken ct = default)
        {
            if (service is null)
            {
                throw new ArgumentNullException(nameof(service));
            }

            if (string.IsNullOrWhiteSpace(token))
            {
                throw new ArgumentException("Token cannot be null or empty.", nameof(token));
            }

            if (rateLimiter is null)
            {
                throw new ArgumentNullException(nameof(rateLimiter));
            }

            if (string.IsNullOrWhiteSpace(subject))
            {
                throw new ArgumentException("Subject cannot be null or empty.", nameof(subject));
            }

            // Check if subject is already rate limited
            var isLocked = await rateLimiter.IsLockedOutAsync(subject).ConfigureAwait(false);
            if (isLocked)
            {
                var timeRemaining = await rateLimiter.GetLockoutTimeRemainingAsync(subject).ConfigureAwait(false);
                return new TokenResult.Malformed(
                    $"Rate limit exceeded. Try again in {timeRemaining?.TotalSeconds:F0} seconds.");
            }

            // Validate the token
            var result = await service.ValidateAsync(token, options, ct).ConfigureAwait(false);

            // If validation failed, record the failure
            if (!result.IsSuccess)
            {
                var recorded = await rateLimiter.TryRecordFailureAsync(subject).ConfigureAwait(false);
                if (!recorded)
                {
                    return new TokenResult.Malformed("Rate limit exceeded for invalid token attempts.");
                }
            }
            else
            {
                // Reset the failure count on successful validation
                await rateLimiter.ResetAsync(subject).ConfigureAwait(false);
            }

            return result;
        }

        /// <summary>
        /// Validates a token and caches the result for subsequent lookups.
        /// Reduces computational overhead for frequently-validated tokens.
        /// </summary>
        /// <param name="service">The token service</param>
        /// <param name="token">Token to validate</param>
        /// <param name="options">Validation options</param>
        /// <param name="cache">The cache to store validated tokens</param>
        /// <param name="ct">Cancellation token</param>
        /// <returns>Validation result</returns>
        public static async Task<TokenResult> ValidateWithCachingAsync(
            this ITokenService service,
            string token,
            TokenValidationOptions options,
            ISecurityTokenCache cache,
            CancellationToken ct = default)
        {
            if (service is null)
            {
                throw new ArgumentNullException(nameof(service));
            }

            if (string.IsNullOrWhiteSpace(token))
            {
                throw new ArgumentException("Token cannot be null or empty.", nameof(token));
            }

            if (cache is null)
            {
                throw new ArgumentNullException(nameof(cache));
            }

            // Try to decode token without verification to get the token ID
            if (!TokenSerializer.TryDecode(token, out var payload, out _))
            {
                return new TokenResult.Malformed("Token format is invalid.");
            }

            var claims = TokenSerializer.Deserialize(payload);
            if (claims is null)
            {
                return new TokenResult.Malformed("Token payload could not be decoded.");
            }

            // Check cache first
            var cached = await cache.TryGetAsync(claims.TokenId).ConfigureAwait(false);
            if (cached)
            {
                return new TokenResult.Success(token, claims);
            }

            // If not in cache, validate normally
            var result = await service.ValidateAsync(token, options, ct).ConfigureAwait(false);

            // If successful, add to cache
            if (result is TokenResult.Success success)
            {
                await cache.SetAsync(success.Claims.TokenId, success.Claims.ExpiresAt).ConfigureAwait(false);
            }

            return result;
        }

        /// <summary>
        /// Validates a token with all available security checks (rate limiting, blacklist, caching).
        /// </summary>
        /// <param name="service">The token service</param>
        /// <param name="token">Token to validate</param>
        /// <param name="options">Validation options</param>
        /// <param name="configuration">Security configuration with rate limiter, blacklist, and/or cache</param>
        /// <param name="rateLimitSubject">Optional subject for rate limiting (IP, user ID, etc.)</param>
        /// <param name="ct">Cancellation token</param>
        /// <returns>Validation result with all security checks applied</returns>
        public static async Task<TokenResult> ValidateWithSecurityAsync(
            this ITokenService service,
            string token,
            TokenValidationOptions options,
            SecureTokenConfiguration configuration,
            string? rateLimitSubject = null,
            CancellationToken ct = default)
        {
            if (service is null)
            {
                throw new ArgumentNullException(nameof(service));
            }

            if (string.IsNullOrWhiteSpace(token))
            {
                throw new ArgumentException("Token cannot be null or empty.", nameof(token));
            }

            if (configuration is null)
            {
                throw new ArgumentNullException(nameof(configuration));
            }

            var result = await service.ValidateAsync(token, options, ct).ConfigureAwait(false);

            // Check blacklist if configured
            if (configuration.Blacklist is not null && result is TokenResult.Success success)
            {
                var isBlacklisted = await configuration.Blacklist.IsBlacklistedAsync(
                    success.Claims.TokenId).ConfigureAwait(false);
                if (isBlacklisted)
                {
                    result = new TokenResult.Revoked(success.Claims.TokenId);
                }
            }

            // Apply rate limiting if configured
            if (configuration.RateLimiter is not null && !string.IsNullOrWhiteSpace(rateLimitSubject))
            {
                if (!result.IsSuccess)
                {
                    var recorded = await configuration.RateLimiter.TryRecordFailureAsync(rateLimitSubject).ConfigureAwait(false);
                    if (!recorded)
                    {
                        return new TokenResult.Malformed("Rate limit exceeded.");
                    }
                }
                else
                {
                    await configuration.RateLimiter.ResetAsync(rateLimitSubject).ConfigureAwait(false);
                }
            }

            // Cache successful validations if configured
            if (configuration.Cache is not null && result is TokenResult.Success cacheSuccess)
            {
                await configuration.Cache.SetAsync(cacheSuccess.Claims.TokenId, cacheSuccess.Claims.ExpiresAt).ConfigureAwait(false);
            }

            return result;
        }

        /// <summary>
        /// Gets a summary of all failed validations for a subject during the current rate limit window.
        /// Useful for security auditing and debugging.
        /// </summary>
        /// <param name="rateLimiter">The rate limiter</param>
        /// <param name="subject">The subject to check</param>
        /// <returns>Summary information including failure count and lockout status</returns>
        public static async Task<string> GetFailureSummaryAsync(
            this ITokenRateLimiter rateLimiter,
            string subject)
        {
            if (rateLimiter is null)
            {
                throw new ArgumentNullException(nameof(rateLimiter));
            }

            if (string.IsNullOrWhiteSpace(subject))
            {
                throw new ArgumentException("Subject cannot be null or empty.", nameof(subject));
            }

            var failureCount = await rateLimiter.GetFailureCountAsync(subject).ConfigureAwait(false);
            var lockoutTime = await rateLimiter.GetLockoutTimeRemainingAsync(subject).ConfigureAwait(false);

            if (lockoutTime.HasValue)
            {
                return $"Subject '{subject}' is locked out. {failureCount} failures recorded. " +
                       $"Time remaining: {lockoutTime.Value.TotalSeconds:F0} seconds.";
            }

            return $"Subject '{subject}' has {failureCount} failures recorded (not locked out).";
        }

        /// <summary>
        /// Introspects and compares multiple tokens for diagnostic purposes.
        /// </summary>
        /// <param name="introspector">The token introspector</param>
        /// <param name="tokens">Tokens to compare</param>
        /// <returns>Formatted comparison of token claims</returns>
        public static string CompareTokens(
            this ITokenIntrospector introspector,
            params string[] tokens)
        {
            if (introspector is null)
            {
                throw new ArgumentNullException(nameof(introspector));
            }

            if (tokens is null || tokens.Length == 0)
            {
                throw new ArgumentException("At least one token must be provided.", nameof(tokens));
            }

            var results = introspector.IntrospectMultiple(tokens);
            var sb = new System.Text.StringBuilder();
            sb.AppendLine("=== Token Comparison ===");

            for (int i = 0; i < results.Count; i++)
            {
                sb.AppendLine($"\n--- Token {i + 1} ---");
                if (results[i].IsValidFormat && results[i].Claims is not null)
                {
                    var claims = results[i].Claims;
                    sb.AppendLine($"Subject: {claims.Subject}");
                    sb.AppendLine($"Issuer: {claims.Issuer}");
                    sb.AppendLine($"Audiences: {string.Join(", ", claims.Audiences)}");
                    sb.AppendLine($"Token Type: {claims.TokenType}");
                    sb.AppendLine($"Expires: {(DateTimeOffset.UtcNow > claims.ExpiresAt ? "EXPIRED" : "VALID")}");
                }
                else
                {
                    sb.AppendLine($"ERROR: {results[i].ErrorMessage}");
                }
            }

            return sb.ToString();
        }
    }
}