tm_back/pkg/authorization/authorization.go

package authorization

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/hex"
	"encoding/pem"
	"fmt"
	"strconv"
	"time"

	"tm/pkg/logger"
	"tm/pkg/redis"

	"github.com/golang-jwt/jwt/v5"
	redisv9 "github.com/redis/go-redis/v9"
)

// TokenType represents the type of token
type TokenType string

const (
	AccessToken  TokenType = "access"
	RefreshToken TokenType = "refresh"
)

// TokenClaims represents the JWT claims structure
type TokenClaims struct {
	UserID    string    `json:"user_id"`
	Email     string    `json:"email"`
	Role      string    `json:"role"`
	CompanyID string    `json:"company_id,omitempty"`
	TokenType TokenType `json:"token_type"`
	jwt.RegisteredClaims
}

// TokenPair represents a pair of access and refresh tokens
type TokenPair struct {
	AccessToken  string `json:"access_token"`
	RefreshToken string `json:"refresh_token"`
	ExpiresIn    int64  `json:"expires_in"`
}

// TokenValidationResult represents the result of token validation
type TokenValidationResult struct {
	Valid   bool
	Claims  *TokenClaims
	Error   string
	Expired bool
}

// AuthorizationService defines the interface for authorization operations
type AuthorizationService interface {
	// GenerateTokenPair generates both access and refresh tokens
	GenerateTokenPair(userID, email, role, companyID string) (*TokenPair, error)

	// GenerateAccessToken generates only an access token
	GenerateAccessToken(userID, email, role, companyID string) (string, error)

	// GenerateRefreshToken generates only a refresh token
	GenerateRefreshToken(userID, email, role, companyID string) (string, error)

	// ValidateToken validates a token and returns the claims
	ValidateToken(tokenString string) (*TokenValidationResult, error)

	// ValidateAccessToken validates specifically an access token
	ValidateAccessToken(tokenString string) (*TokenValidationResult, error)

	// ValidateRefreshToken validates specifically a refresh token
	ValidateRefreshToken(tokenString string) (*TokenValidationResult, error)

	// RefreshAccessToken generates a new access token using a valid refresh token
	RefreshAccessToken(refreshToken string) (string, error)

	// InvalidateToken adds a token to the blacklist (if using Redis)
	InvalidateToken(tokenString string) error

	// IsTokenBlacklisted checks if a token is blacklisted
	IsTokenBlacklisted(tokenString string) (bool, error)

	// ClearExpiredTokens removes expired tokens from the blacklist
	ClearExpiredTokens() error

	// GetBlacklistStats returns statistics about the blacklist
	GetBlacklistStats(ctx context.Context) (map[string]interface{}, error)

	// InvalidateUserSessions marks all tokens issued before now as invalid for the user.
	InvalidateUserSessions(userID string) error
}

// AuthorizationConfig holds the configuration for the authorization service
type AuthorizationConfig struct {
	AccessTokenSecret  string        `json:"access_token_secret"`
	RefreshTokenSecret string        `json:"refresh_token_secret"`
	AccessTokenExpiry  time.Duration `json:"access_token_expiry"`
	RefreshTokenExpiry time.Duration `json:"refresh_token_expiry"`
	Issuer             string        `json:"issuer"`
	Audience           string        `json:"audience"`
}

// DefaultConfig returns a default configuration
func DefaultConfig() *AuthorizationConfig {
	return &AuthorizationConfig{
		AccessTokenSecret:  "your-access-token-secret-key-here",
		RefreshTokenSecret: "your-refresh-token-secret-key-here",
		AccessTokenExpiry:  15 * time.Minute,   // 15 minutes
		RefreshTokenExpiry: 7 * 24 * time.Hour, // 7 days
		Issuer:             "tender-management",
		Audience:           "tender-management-api",
	}
}

// Service implements the AuthorizationService interface
type Service struct {
	config *AuthorizationConfig
	logger logger.Logger
	redis  redis.Client
}

// NewAuthorizationService creates a new authorization service
func NewAuthorizationService(config *AuthorizationConfig, logger logger.Logger, redis redis.Client) AuthorizationService {
	if config == nil {
		config = DefaultConfig()
	}

	return &Service{
		config: config,
		logger: logger,
		redis:  redis,
	}
}

// GenerateTokenPair generates both access and refresh tokens
func (s *Service) GenerateTokenPair(userID, email, role, companyID string) (*TokenPair, error) {
	s.logger.Info("Generating token pair", map[string]interface{}{
		"user_id": userID,
		"email":   email,
		"role":    role,
	})

	accessToken, err := s.GenerateAccessToken(userID, email, role, companyID)
	if err != nil {
		s.logger.Error("Failed to generate access token", map[string]interface{}{
			"user_id": userID,
			"error":   err.Error(),
		})
		return nil, fmt.Errorf("failed to generate access token: %w", err)
	}

	refreshToken, err := s.GenerateRefreshToken(userID, email, role, companyID)
	if err != nil {
		s.logger.Error("Failed to generate refresh token", map[string]interface{}{
			"user_id": userID,
			"error":   err.Error(),
		})
		return nil, fmt.Errorf("failed to generate refresh token: %w", err)
	}

	tokenPair := &TokenPair{
		AccessToken:  accessToken,
		RefreshToken: refreshToken,
		ExpiresIn:    int64(s.config.AccessTokenExpiry.Seconds()),
	}

	s.logger.Info("Token pair generated successfully", map[string]interface{}{
		"user_id":    userID,
		"expires_in": tokenPair.ExpiresIn,
	})

	return tokenPair, nil
}

// GenerateAccessToken generates only an access token
func (s *Service) GenerateAccessToken(userID, email, role, companyID string) (string, error) {
	now := time.Now()
	expiresAt := now.Add(s.config.AccessTokenExpiry)

	claims := &TokenClaims{
		UserID:    userID,
		Email:     email,
		Role:      role,
		CompanyID: companyID,
		TokenType: AccessToken,
		RegisteredClaims: jwt.RegisteredClaims{
			Issuer:    s.config.Issuer,
			Subject:   userID,
			Audience:  []string{s.config.Audience},
			ExpiresAt: jwt.NewNumericDate(expiresAt),
			NotBefore: jwt.NewNumericDate(now),
			IssuedAt:  jwt.NewNumericDate(now),
		},
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	tokenString, err := token.SignedString([]byte(s.config.AccessTokenSecret))
	if err != nil {
		return "", fmt.Errorf("failed to sign access token: %w", err)
	}

	return tokenString, nil
}

// GenerateRefreshToken generates only a refresh token
func (s *Service) GenerateRefreshToken(userID, email, role, companyID string) (string, error) {
	now := time.Now()
	expiresAt := now.Add(s.config.RefreshTokenExpiry)

	claims := &TokenClaims{
		UserID:    userID,
		Email:     email,
		Role:      role,
		CompanyID: companyID,
		TokenType: RefreshToken,
		RegisteredClaims: jwt.RegisteredClaims{
			Issuer:    s.config.Issuer,
			Subject:   userID,
			Audience:  []string{s.config.Audience},
			ExpiresAt: jwt.NewNumericDate(expiresAt),
			NotBefore: jwt.NewNumericDate(now),
			IssuedAt:  jwt.NewNumericDate(now),
		},
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	tokenString, err := token.SignedString([]byte(s.config.RefreshTokenSecret))
	if err != nil {
		return "", fmt.Errorf("failed to sign refresh token: %w", err)
	}

	return tokenString, nil
}

// ValidateToken validates a token and returns the claims
func (s *Service) ValidateToken(tokenString string) (*TokenValidationResult, error) {
	// First check if token is blacklisted
	isBlacklisted, err := s.IsTokenBlacklisted(tokenString)
	if err != nil {
		s.logger.Error("Failed to check token blacklist", map[string]interface{}{
			"error": err.Error(),
		})
		// Continue with validation even if blacklist check fails
	}

	if isBlacklisted {
		return &TokenValidationResult{
			Valid:   false,
			Claims:  nil,
			Error:   "token is blacklisted",
			Expired: false,
		}, nil
	}

	// Parse the token
	token, err := jwt.ParseWithClaims(tokenString, &TokenClaims{}, func(token *jwt.Token) (interface{}, error) {
		// Validate the signing method
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
		}

		// Determine which secret to use based on token type
		claims, ok := token.Claims.(*TokenClaims)
		if !ok {
			return nil, fmt.Errorf("invalid token claims")
		}

		switch claims.TokenType {
		case AccessToken:
			return []byte(s.config.AccessTokenSecret), nil
		case RefreshToken:
			return []byte(s.config.RefreshTokenSecret), nil
		default:
			return nil, fmt.Errorf("invalid token type")
		}
	})

	if err != nil {
		if jwt.ErrTokenExpired.Error() == err.Error() {
			return &TokenValidationResult{
				Valid:   false,
				Claims:  nil,
				Error:   "token expired",
				Expired: true,
			}, nil
		}

		return &TokenValidationResult{
			Valid:   false,
			Claims:  nil,
			Error:   fmt.Sprintf("token validation failed: %s", err.Error()),
			Expired: false,
		}, nil
	}

	if !token.Valid {
		return &TokenValidationResult{
			Valid:   false,
			Claims:  nil,
			Error:   "invalid token",
			Expired: false,
		}, nil
	}

	claims, ok := token.Claims.(*TokenClaims)
	if !ok {
		return &TokenValidationResult{
			Valid:   false,
			Claims:  nil,
			Error:   "invalid token claims",
			Expired: false,
		}, nil
	}

	if invalidated, err := s.isSessionInvalidated(claims.UserID, claims.IssuedAt); err != nil {
		s.logger.Error("Failed to check session invalidation", map[string]interface{}{
			"user_id": claims.UserID,
			"error":   err.Error(),
		})
	} else if invalidated {
		return &TokenValidationResult{
			Valid:   false,
			Claims:  nil,
			Error:   "session invalidated",
			Expired: false,
		}, nil
	}

	return &TokenValidationResult{
		Valid:   true,
		Claims:  claims,
		Error:   "",
		Expired: false,
	}, nil
}

// ValidateAccessToken validates specifically an access token
func (s *Service) ValidateAccessToken(tokenString string) (*TokenValidationResult, error) {
	result, err := s.ValidateToken(tokenString)
	if err != nil {
		return nil, err
	}

	if !result.Valid {
		return result, nil
	}

	// Check if it's actually an access token
	if result.Claims.TokenType != AccessToken {
		return &TokenValidationResult{
			Valid:   false,
			Claims:  nil,
			Error:   "token is not an access token",
			Expired: false,
		}, nil
	}

	return result, nil
}

// ValidateRefreshToken validates specifically a refresh token
func (s *Service) ValidateRefreshToken(tokenString string) (*TokenValidationResult, error) {
	result, err := s.ValidateToken(tokenString)
	if err != nil {
		return nil, err
	}

	if !result.Valid {
		return result, nil
	}

	// Check if it's actually a refresh token
	if result.Claims.TokenType != RefreshToken {
		return &TokenValidationResult{
			Valid:   false,
			Claims:  nil,
			Error:   "token is not a refresh token",
			Expired: false,
		}, nil
	}

	return result, nil
}

// RefreshAccessToken generates a new access token using a valid refresh token
func (s *Service) RefreshAccessToken(refreshToken string) (string, error) {
	s.logger.Info("Refreshing access token", map[string]interface{}{
		"refresh_token_length": len(refreshToken),
	})

	// Validate the refresh token
	result, err := s.ValidateRefreshToken(refreshToken)
	if err != nil {
		s.logger.Error("Failed to validate refresh token", map[string]interface{}{
			"error": err.Error(),
		})
		return "", fmt.Errorf("failed to validate refresh token: %w", err)
	}

	if !result.Valid {
		s.logger.Warn("Invalid refresh token provided", map[string]interface{}{
			"error": result.Error,
		})
		return "", fmt.Errorf("invalid refresh token: %s", result.Error)
	}

	// Generate new access token
	newAccessToken, err := s.GenerateAccessToken(
		result.Claims.UserID,
		result.Claims.Email,
		result.Claims.Role,
		result.Claims.CompanyID,
	)
	if err != nil {
		s.logger.Error("Failed to generate new access token", map[string]interface{}{
			"user_id": result.Claims.UserID,
			"error":   err.Error(),
		})
		return "", fmt.Errorf("failed to generate new access token: %w", err)
	}

	s.logger.Info("Access token refreshed successfully", map[string]interface{}{
		"user_id": result.Claims.UserID,
	})

	return newAccessToken, nil
}

func sessionInvalidationKey(userID string) string {
	return "sessions:invalidated_at:" + userID
}

// InvalidateUserSessions marks tokens issued before now as invalid for the given user.
func (s *Service) InvalidateUserSessions(userID string) error {
	if s.redis == nil {
		s.logger.Warn("Redis client not available, session invalidation skipped", map[string]interface{}{
			"user_id": userID,
		})
		return nil
	}

	now := time.Now().Unix()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	ttl := s.config.RefreshTokenExpiry + 5*time.Minute
	if err := s.redis.Set(ctx, sessionInvalidationKey(userID), strconv.FormatInt(now, 10), ttl); err != nil {
		return fmt.Errorf("failed to invalidate user sessions: %w", err)
	}

	s.logger.Info("User sessions invalidated", map[string]interface{}{
		"user_id": userID,
	})
	return nil
}

func (s *Service) isSessionInvalidated(userID string, issuedAt *jwt.NumericDate) (bool, error) {
	if s.redis == nil || issuedAt == nil {
		return false, nil
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	value, err := s.redis.Get(ctx, sessionInvalidationKey(userID))
	if err != nil {
		if err == redisv9.Nil {
			return false, nil
		}
		return false, err
	}

	invalidatedAt, err := strconv.ParseInt(value, 10, 64)
	if err != nil {
		return false, err
	}

	return issuedAt.Time.Unix() < invalidatedAt, nil
}

// generateTokenHash generates a SHA-256 hash of the token for storage
func (s *Service) generateTokenHash(tokenString string) string {
	hash := sha256.Sum256([]byte(tokenString))
	return hex.EncodeToString(hash[:])
}

// InvalidateToken adds a token to the blacklist
func (s *Service) InvalidateToken(tokenString string) error {
	if s.redis == nil {
		s.logger.Warn("Redis client not available, token blacklisting disabled", map[string]interface{}{
			"token_length": len(tokenString),
		})
		return nil
	}

	// Generate hash of the token for storage
	tokenHash := s.generateTokenHash(tokenString)

	// Parse token to get expiration time for blacklist TTL
	claims := &TokenClaims{}
	_, err := jwt.ParseWithClaims(tokenString, claims, func(token *jwt.Token) (interface{}, error) {
		// Determine which secret to use based on token type
		claims, ok := token.Claims.(*TokenClaims)
		if !ok {
			return nil, fmt.Errorf("invalid token claims")
		}

		switch claims.TokenType {
		case AccessToken:
			return []byte(s.config.AccessTokenSecret), nil
		case RefreshToken:
			return []byte(s.config.RefreshTokenSecret), nil
		default:
			return nil, fmt.Errorf("invalid token type")
		}
	})

	if err != nil {
		s.logger.Error("Failed to parse token for blacklisting", map[string]interface{}{
			"error": err.Error(),
		})
		// Still add to blacklist with default TTL even if parsing fails
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()

		// Use default TTL (24 hours) if we can't determine actual expiration
		defaultTTL := 24 * time.Hour
		err = s.redis.Set(ctx, "blacklist:"+tokenHash, "1", defaultTTL)
		if err != nil {
			s.logger.Error("Failed to add token to blacklist", map[string]interface{}{
				"error": err.Error(),
			})
			return fmt.Errorf("failed to add token to blacklist: %w", err)
		}

		s.logger.Info("Token added to blacklist with default TTL", map[string]interface{}{
			"token_hash": tokenHash,
			"ttl":        defaultTTL.String(),
		})
		return nil
	}

	// Calculate TTL based on token expiration
	var ttl time.Duration
	if claims.ExpiresAt != nil {
		ttl = time.Until(claims.ExpiresAt.Time)
		// Add some buffer time to ensure token is blacklisted until it would have expired
		ttl += 5 * time.Minute
	} else {
		// Fallback to default TTL if no expiration
		ttl = 24 * time.Hour
	}

	// Add token to blacklist with TTL
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	err = s.redis.Set(ctx, "blacklist:"+tokenHash, "1", ttl)
	if err != nil {
		s.logger.Error("Failed to add token to blacklist", map[string]interface{}{
			"token_hash": tokenHash,
			"error":      err.Error(),
		})
		return fmt.Errorf("failed to add token to blacklist: %w", err)
	}

	s.logger.Info("Token successfully added to blacklist", map[string]interface{}{
		"token_hash": tokenHash,
		"ttl":        ttl.String(),
	})

	return nil
}

// IsTokenBlacklisted checks if a token is blacklisted
func (s *Service) IsTokenBlacklisted(tokenString string) (bool, error) {
	if s.redis == nil {
		// If Redis is not available, assume token is not blacklisted
		// This allows the system to continue functioning without Redis
		return false, nil
	}

	// Generate hash of the token for lookup
	tokenHash := s.generateTokenHash(tokenString)

	// Check if token exists in blacklist
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	exists, err := s.redis.Exists(ctx, "blacklist:"+tokenHash)
	if err != nil {
		s.logger.Error("Failed to check token blacklist", map[string]interface{}{
			"token_hash": tokenHash,
			"error":      err.Error(),
		})
		// If we can't check the blacklist, assume token is not blacklisted
		// This prevents blocking legitimate requests due to Redis issues
		return false, nil
	}

	isBlacklisted := exists > 0

	if isBlacklisted {
		s.logger.Debug("Token found in blacklist", map[string]interface{}{
			"token_hash": tokenHash,
		})
	} else {
		s.logger.Debug("Token not found in blacklist", map[string]interface{}{
			"token_hash": tokenHash,
		})
	}

	return isBlacklisted, nil
}

// ClearExpiredTokens removes expired tokens from the blacklist
// This is useful for cleanup operations
func (s *Service) ClearExpiredTokens() error {
	if s.redis == nil {
		return nil
	}

	s.logger.Info("Clearing expired tokens from blacklist", map[string]interface{}{})

	// Note: Redis automatically expires keys based on TTL
	// This method is provided for manual cleanup if needed
	// In most cases, Redis TTL handling is sufficient

	return nil
}

// GetBlacklistStats returns statistics about the blacklist
func (s *Service) GetBlacklistStats(ctx context.Context) (map[string]interface{}, error) {
	if s.redis == nil {
		return map[string]interface{}{
			"redis_available": false,
			"message":         "Redis client not available",
		}, nil
	}

	// This is a placeholder for future implementation
	// Could return metrics like total blacklisted tokens, memory usage, etc.
	return map[string]interface{}{
		"redis_available": true,
		"message":         "Blacklist statistics not yet implemented",
	}, nil
}

// GenerateRandomSecret generates a random secret key for JWT signing
func GenerateRandomSecret(length int) (string, error) {
	if length < 32 {
		length = 32 // Minimum recommended length
	}

	bytes := make([]byte, length)
	_, err := rand.Read(bytes)
	if err != nil {
		return "", fmt.Errorf("failed to generate random bytes: %w", err)
	}

	return fmt.Sprintf("%x", bytes), nil
}

// GenerateRSAKeyPair generates an RSA key pair for JWT signing
func GenerateRSAKeyPair(bits int) (privateKeyPEM, publicKeyPEM string, err error) {
	if bits < 2048 {
		bits = 2048 // Minimum recommended bits
	}

	privateKey, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		return "", "", fmt.Errorf("failed to generate RSA key: %w", err)
	}

	// Encode private key
	privateKeyBytes := x509.MarshalPKCS1PrivateKey(privateKey)
	privateKeyPEM = string(pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: privateKeyBytes,
	}))

	// Encode public key
	publicKeyBytes := x509.MarshalPKCS1PublicKey(&privateKey.PublicKey)
	publicKeyPEM = string(pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PUBLIC KEY",
		Bytes: publicKeyBytes,
	}))

	return privateKeyPEM, publicKeyPEM, nil
}