API Key Management

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95

import secrets
import hashlib
import base64
from datetime import datetime, timezone
from typing import Tuple
 
class SecureAPIKeyGenerator:
    """
    Production-grade API key generator with proper entropy,
    prefix support, and checksum validation.
    """
    
    # Prefix identifies key type and enables quick routing
    PREFIX_LIVE = "sk_live_"   # Production keys
    PREFIX_TEST = "sk_test_"  # Testing/sandbox keys
    PREFIX_RESTRICTED = "rk_"  # Restricted scope keys
    
    KEY_BYTES = 32  # 256 bits of entropy
    
    @classmethod
    def generate_key(cls, environment: str = "live") -> Tuple[str, str]:
        """
        Generate a cryptographically secure API key.
        
        Returns:
            Tuple of (full_key, key_hash) where:
            - full_key: The complete API key to give to the client
            - key_hash: SHA-256 hash to store in database
        """
        # Use secrets module (CSPRNG) - NEVER use random module
        raw_bytes = secrets.token_bytes(cls.KEY_BYTES)
        
        # URL-safe base64 encoding for easy transmission
        key_body = base64.urlsafe_b64encode(raw_bytes).decode('utf-8')
        
        # Remove padding for cleaner keys
        key_body = key_body.rstrip('=')
        
        # Add environment-specific prefix
        prefix = cls.PREFIX_LIVE if environment == "live" else cls.PREFIX_TEST
        
        # Add 4-character checksum for typo detection
        checksum = cls._compute_checksum(key_body)
        
        full_key = f"{prefix}{key_body}_{checksum}"
        
        # CRITICAL: Never store the raw key, only its hash
        key_hash = hashlib.sha256(full_key.encode()).hexdigest()
        
        return full_key, key_hash
    
    @classmethod
    def _compute_checksum(cls, key_body: str) -> str:
        """Compute 4-character checksum for typo detection."""
        hash_bytes = hashlib.sha256(key_body.encode()).digest()
        return base64.urlsafe_b64encode(hash_bytes[:3]).decode('utf-8')[:4]
    
    @classmethod
    def validate_format(cls, api_key: str) -> bool:
        """
        Validate API key format WITHOUT checking database.
        Useful for quick rejection of malformed keys.
        """
        # Check prefix
        valid_prefixes = (cls.PREFIX_LIVE, cls.PREFIX_TEST, cls.PREFIX_RESTRICTED)
        if not any(api_key.startswith(p) for p in valid_prefixes):
            return False
        
        # Extract parts
        parts = api_key.split('_')
        if len(parts) < 3:
            return False
            
        # Verify checksum
        key_body = '_'.join(parts[2:-1]) if len(parts) > 3 else parts[2].rsplit('_', 1)[0]
        provided_checksum = parts[-1]
        expected_checksum = cls._compute_checksum(key_body)
        
        return secrets.compare_digest(provided_checksum, expected_checksum)
    
    @classmethod
    def hash_key_for_lookup(cls, api_key: str) -> str:
        """Hash the key for secure database lookup."""
        return hashlib.sha256(api_key.encode()).hexdigest()
 
 
# Example usage
if __name__ == "__main__":
    generator = SecureAPIKeyGenerator()
    
    # Generate production key
    api_key, key_hash = generator.generate_key("live")
    print(f"API Key (give to client): {api_key}")
    print(f"Key Hash (store in DB):   {key_hash}")
    print(f"Format Valid: {generator.validate_format(api_key)}")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98

-- API Keys table with security-first design
CREATE TABLE api_keys (
    -- Internal identifier (never exposed)
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    
    -- SHA-256 hash of the full API key (for lookup)
    key_hash VARCHAR(64) UNIQUE NOT NULL,
    
    -- Key prefix for display (e.g., "sk_live_abc...")
    -- Only first 12 chars stored for identification
    key_prefix VARCHAR(20) NOT NULL,
    
    -- Owner information
    organization_id UUID NOT NULL REFERENCES organizations(id),
    created_by_user_id UUID REFERENCES users(id),
    
    -- Human-readable name for the key
    name VARCHAR(255) NOT NULL,
    
    -- Key metadata
    environment VARCHAR(20) NOT NULL CHECK (environment IN ('live', 'test')),
    
    -- Permissions (JSON array of allowed scopes)
    scopes JSONB NOT NULL DEFAULT '[]',
    
    -- Rate limiting configuration
    rate_limit_per_minute INTEGER DEFAULT 1000,
    rate_limit_per_day INTEGER DEFAULT 100000,
    
    -- IP restrictions (NULL = allow all)
    allowed_ips INET[] DEFAULT NULL,
    allowed_cidrs CIDR[] DEFAULT NULL,
    
    -- Lifecycle management
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    last_used_at TIMESTAMPTZ,
    expires_at TIMESTAMPTZ,  -- NULL = never expires
    revoked_at TIMESTAMPTZ,
    revoked_by_user_id UUID REFERENCES users(id),
    revocation_reason TEXT,
    
    -- Key version for rotation tracking
    version INTEGER NOT NULL DEFAULT 1,
    
    -- Indexes for performance
    CONSTRAINT valid_expiration CHECK (expires_at IS NULL OR expires_at > created_at)
);
 
-- Index for fast key lookup by hash (most critical index)
CREATE UNIQUE INDEX idx_api_keys_hash ON api_keys(key_hash) 
WHERE revoked_at IS NULL;
 
-- Index for listing keys by organization
CREATE INDEX idx_api_keys_org ON api_keys(organization_id, created_at DESC);
 
-- Index for finding expired keys (cleanup job)
CREATE INDEX idx_api_keys_expires 
ON api_keys(expires_at) 
WHERE expires_at IS NOT NULL AND revoked_at IS NULL;
 
-- Index for usage tracking updates
CREATE INDEX idx_api_keys_last_used 
ON api_keys(last_used_at) 
WHERE revoked_at IS NULL;
 
-- ================================================
-- API Key Usage Log (for analytics and auditing)
-- ================================================
CREATE TABLE api_key_usage_log (
    id BIGSERIAL PRIMARY KEY,
    key_id UUID NOT NULL REFERENCES api_keys(id),
    
    -- Request metadata
    timestamp TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    endpoint VARCHAR(500) NOT NULL,
    method VARCHAR(10) NOT NULL,
    status_code INTEGER NOT NULL,
    response_time_ms INTEGER,
    
    -- Client information
    client_ip INET NOT NULL,
    user_agent TEXT,
    
    -- Request identifiers
    request_id UUID,
    
    -- Partition by month for efficient cleanup
    created_month DATE NOT NULL DEFAULT date_trunc('month', NOW())
) PARTITION BY RANGE (created_month);
 
-- Create partitions (example for current and next month)
CREATE TABLE api_key_usage_log_2024_01 
PARTITION OF api_key_usage_log 
FOR VALUES FROM ('2024-01-01') TO ('2024-02-01');
 
-- Index for querying by key
CREATE INDEX idx_usage_log_key_time 
ON api_key_usage_log(key_id, timestamp DESC);

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155

from datetime import datetime, timedelta, timezone
from typing import Optional, Tuple
from dataclasses import dataclass
import logging
 
logger = logging.getLogger(__name__)
 
@dataclass
class APIKey:
    id: str
    key_hash: str
    organization_id: str
    version: int
    created_at: datetime
    expires_at: Optional[datetime]
    revoked_at: Optional[datetime]
    
class KeyRotationService:
    """
    Implements zero-downtime API key rotation with grace periods.
    """
    
    # How long both keys are valid during rotation
    GRACE_PERIOD_DAYS = 7
    
    # Warning threshold before expiration
    WARNING_THRESHOLD_DAYS = 14
    
    def __init__(self, key_repository, key_generator, notification_service):
        self.key_repository = key_repository
        self.key_generator = key_generator
        self.notification_service = notification_service
    
    def rotate_key(
        self, 
        current_key_id: str,
        initiated_by: str,
        reason: str = "scheduled_rotation"
    ) -> Tuple[str, datetime]:
        """
        Rotate an API key with grace period for zero-downtime transition.
        
        Returns:
            Tuple of (new_api_key, old_key_expiration)
        """
        # Fetch current key
        current_key = self.key_repository.get_by_id(current_key_id)
        if not current_key:
            raise ValueError(f"Key {current_key_id} not found")
        
        if current_key.revoked_at:
            raise ValueError("Cannot rotate a revoked key")
        
        # Generate new key with incremented version
        new_api_key, new_key_hash = self.key_generator.generate_key(
            environment=current_key.environment
        )
        
        # Calculate when old key should expire
        old_key_expiration = datetime.now(timezone.utc) + timedelta(
            days=self.GRACE_PERIOD_DAYS
        )
        
        # Create new key entry
        new_key_record = self.key_repository.create(
            key_hash=new_key_hash,
            key_prefix=new_api_key[:20] + "...",
            organization_id=current_key.organization_id,
            version=current_key.version + 1,
            scopes=current_key.scopes,  # Inherit permissions
            rate_limits=current_key.rate_limits,  # Inherit limits
            name=f"{current_key.name} (rotated)",
            previous_key_id=current_key.id,
        )
        
        # Schedule old key expiration (soft delete)
        self.key_repository.schedule_expiration(
            key_id=current_key.id,
            expires_at=old_key_expiration
        )
        
        # Log the rotation event
        self.key_repository.log_event(
            key_id=current_key.id,
            event_type="rotation_initiated",
            initiated_by=initiated_by,
            reason=reason,
            new_key_id=new_key_record.id,
            old_key_expires_at=old_key_expiration,
        )
        
        # Notify organization admins
        self.notification_service.notify_key_rotation(
            organization_id=current_key.organization_id,
            old_key_prefix=current_key.key_prefix,
            old_key_expires_at=old_key_expiration,
            reason=reason,
        )
        
        logger.info(
            f"Rotated key {current_key.id} -> {new_key_record.id}, "
            f"old key expires at {old_key_expiration}"
        )
        
        return new_api_key, old_key_expiration
    
    def check_keys_needing_rotation(self) -> list:
        """
        Find keys approaching expiration that need rotation warnings.
        Run this as a scheduled job.
        """
        warning_threshold = datetime.now(timezone.utc) + timedelta(
            days=self.WARNING_THRESHOLD_DAYS
        )
        
        keys = self.key_repository.find_expiring_before(warning_threshold)
        
        for key in keys:
            self.notification_service.notify_key_expiring(
                organization_id=key.organization_id,
                key_prefix=key.key_prefix,
                expires_at=key.expires_at,
                days_remaining=(key.expires_at - datetime.now(timezone.utc)).days
            )
        
        return keys
    
    def emergency_revoke_all(
        self, 
        organization_id: str,
        initiated_by: str,
        reason: str
    ) -> int:
        """
        Emergency revocation of all keys for an organization.
        Use in case of suspected breach.
        """
        revoked_count = self.key_repository.revoke_all_for_org(
            organization_id=organization_id,
            revoked_by=initiated_by,
            reason=reason,
        )
        
        self.notification_service.notify_emergency_revocation(
            organization_id=organization_id,
            revoked_count=revoked_count,
            reason=reason,
        )
        
        logger.critical(
            f"EMERGENCY: Revoked {revoked_count} keys for org {organization_id}. "
            f"Reason: {reason}"
        )
        
        return revoked_count

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173

package apikey
 
import (
    "context"
    "encoding/json"
    "fmt"
    "time"
    
    "github.com/redis/go-redis/v9"
)
 
// RevocationService handles API key revocation with cache invalidation
type RevocationService struct {
    db          *KeyRepository
    redis       *redis.Client
    pubsub      *redis.PubSub
    auditLogger *AuditLogger
}
 
// RevocationEvent is published when a key is revoked
type RevocationEvent struct {
    KeyHash   string    `json:"key_hash"`
    KeyID     string    `json:"key_id"`
    OrgID     string    `json:"org_id"`
    RevokedAt time.Time `json:"revoked_at"`
    RevokedBy string    `json:"revoked_by"`
    Reason    string    `json:"reason"`
}
 
const (
    revocationChannel = "api_key_revocations"
    keyCacheTTL       = 60 * time.Second
    keyPrefix         = "apikey:"
)
 
// RevokeKey immediately revokes an API key and broadcasts invalidation
func (s *RevocationService) RevokeKey(
    ctx context.Context,
    keyID string,
    revokedBy string,
    reason string,
) error {
    // 1. Update database (source of truth)
    key, err := s.db.RevokeKey(ctx, keyID, revokedBy, reason)
    if err != nil {
        return fmt.Errorf("failed to revoke key in database: %w", err)
    }
    
    // 2. Immediately delete from local cache
    cacheKey := keyPrefix + key.KeyHash
    if err := s.redis.Del(ctx, cacheKey).Err(); err != nil {
        // Log but don't fail - cache will expire
        s.auditLogger.Warn("cache delete failed", "key_id", keyID, "error", err)
    }
    
    // 3. Broadcast revocation to all instances
    event := RevocationEvent{
        KeyHash:   key.KeyHash,
        KeyID:     key.ID,
        OrgID:     key.OrganizationID,
        RevokedAt: time.Now().UTC(),
        RevokedBy: revokedBy,
        Reason:    reason,
    }
    
    eventJSON, err := json.Marshal(event)
    if err != nil {
        return fmt.Errorf("failed to marshal revocation event: %w", err)
    }
    
    if err := s.redis.Publish(ctx, revocationChannel, eventJSON).Err(); err != nil {
        // Log critical - other instances won't know about revocation
        s.auditLogger.Error("CRITICAL: revocation broadcast failed", 
            "key_id", keyID, 
            "error", err,
        )
        // Don't return error - revocation succeeded in DB
    }
    
    // 4. Log to audit trail
    s.auditLogger.Log(AuditEvent{
        Type:      "key_revoked",
        KeyID:     keyID,
        Actor:     revokedBy,
        Reason:    reason,
        Timestamp: event.RevokedAt,
    })
    
    return nil
}
 
// StartRevocationListener subscribes to revocation events
func (s *RevocationService) StartRevocationListener(ctx context.Context) {
    s.pubsub = s.redis.Subscribe(ctx, revocationChannel)
    
    go func() {
        ch := s.pubsub.Channel()
        for {
            select {
            case <-ctx.Done():
                s.pubsub.Close()
                return
            case msg := <-ch:
                var event RevocationEvent
                if err := json.Unmarshal([]byte(msg.Payload), &event); err != nil {
                    s.auditLogger.Error("failed to unmarshal revocation event", 
                        "error", err,
                    )
                    continue
                }
                
                // Invalidate local cache
                cacheKey := keyPrefix + event.KeyHash
                s.redis.Del(ctx, cacheKey)
                
                s.auditLogger.Info("processed revocation broadcast",
                    "key_id", event.KeyID,
                )
            }
        }
    }()
}
 
// ValidateKey checks if a key is valid, using cache with revocation awareness
func (s *RevocationService) ValidateKey(
    ctx context.Context,
    keyHash string,
) (*KeyMetadata, error) {
    cacheKey := keyPrefix + keyHash
    
    // Try cache first
    cached, err := s.redis.Get(ctx, cacheKey).Bytes()
    if err == nil {
        var metadata KeyMetadata
        if err := json.Unmarshal(cached, &metadata); err == nil {
            // Check if cached entry is marked as revoked
            if metadata.RevokedAt != nil {
                return nil, ErrKeyRevoked
            }
            return &metadata, nil
        }
    }
    
    // Cache miss - query database
    key, err := s.db.GetKeyByHash(ctx, keyHash)
    if err != nil {
        return nil, err
    }
    
    if key == nil {
        // Cache negative result briefly to prevent repeated lookups
        s.redis.Set(ctx, cacheKey, []byte("{}"), 30*time.Second)
        return nil, ErrKeyNotFound
    }
    
    if key.RevokedAt != nil {
        return nil, ErrKeyRevoked
    }
    
    // Cache valid key
    metadata := &KeyMetadata{
        KeyID:        key.ID,
        OrgID:        key.OrganizationID,
        Scopes:       key.Scopes,
        RateLimits:   key.RateLimits,
        RevokedAt:    key.RevokedAt,
    }
    
    cacheBytes, _ := json.Marshal(metadata)
    s.redis.Set(ctx, cacheKey, cacheBytes, keyCacheTTL)
    
    return metadata, nil
}