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

import os
import secrets
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives.asymmetric import rsa, ec
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
 
# Method 1: Generate raw symmetric key (AES-256)
# Uses OS-level CSPRNG (Cryptographically Secure PRNG)
symmetric_key = secrets.token_bytes(32)  # 256 bits
print(f"AES-256 key: {symmetric_key.hex()}")
 
# Method 2: Generate AES-GCM key using cryptography library
# This is the preferred method for authenticated encryption
aesgcm_key = AESGCM.generate_key(bit_length=256)
 
# Method 3: Generate RSA key pair
# 4096 bits recommended for long-term security
rsa_private_key = rsa.generate_private_key(
    public_exponent=65537,  # Standard secure exponent
    key_size=4096,          # Bit length
)
rsa_public_key = rsa_private_key.public_key()
 
# Method 4: Generate ECDSA key pair (more efficient than RSA)
# P-384 provides ~192-bit security level
ec_private_key = ec.generate_private_key(ec.SECP384R1())
ec_public_key = ec_private_key.public_key()
 
# Method 5: Derive key from password (when you must use passwords)
# Always use proper KDF with high iteration count
password = b"user_provided_password"
salt = os.urandom(16)  # Random salt per password
kdf = PBKDF2HMAC(
    algorithm=hashes.SHA256(),
    length=32,           # 256-bit derived key
    salt=salt,
    iterations=600000,   # OWASP 2023 recommendation
)
derived_key = kdf.derive(password)
 
# NEVER do this:
# weak_key = hashlib.md5(password).digest()  # No salt, fast hash
# weak_key = str(time.time()).encode()[:32]  # Predictable

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

import boto3
import base64
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
import os
 
class EnvelopeEncryption:
    """
    Demonstrates envelope encryption using AWS KMS as the KEK provider.
    The KEK never leaves AWS KMS; we only receive wrapped/unwrapped DEKs.
    """
    
    def __init__(self, kms_key_id: str):
        self.kms = boto3.client('kms')
        self.kms_key_id = kms_key_id
    
    def encrypt_data(self, plaintext: bytes) -> dict:
        """
        Encrypt data using envelope encryption.
        
        1. Request a new data key from KMS
        2. KMS returns plaintext DEK + encrypted DEK
        3. Use plaintext DEK to encrypt data locally
        4. Discard plaintext DEK
        5. Store encrypted data + encrypted DEK together
        """
        # Generate data key - KMS returns both plaintext and encrypted versions
        data_key_response = self.kms.generate_data_key(
            KeyId=self.kms_key_id,
            KeySpec='AES_256'  # 256-bit AES key
        )
        
        plaintext_dek = data_key_response['Plaintext']     # 32 bytes
        encrypted_dek = data_key_response['CiphertextBlob'] # Wrapped key
        
        # Use the plaintext DEK for local encryption
        nonce = os.urandom(12)  # 96-bit nonce for AES-GCM
        aesgcm = AESGCM(plaintext_dek)
        ciphertext = aesgcm.encrypt(nonce, plaintext, None)
        
        # CRITICAL: Zero out the plaintext key from memory
        # Python doesn't guarantee this, but we try
        plaintext_dek = b'\x00' * len(plaintext_dek)
        del plaintext_dek
        
        # Return envelope: encrypted data + encrypted key + nonce
        return {
            'encrypted_data': base64.b64encode(ciphertext).decode(),
            'encrypted_dek': base64.b64encode(encrypted_dek).decode(),
            'nonce': base64.b64encode(nonce).decode(),
        }
    
    def decrypt_data(self, envelope: dict) -> bytes:
        """
        Decrypt data using envelope encryption.
        
        1. Extract encrypted DEK from envelope
        2. Ask KMS to decrypt (unwrap) the DEK
        3. Use plaintext DEK to decrypt data locally
        4. Discard plaintext DEK
        """
        encrypted_dek = base64.b64decode(envelope['encrypted_dek'])
        ciphertext = base64.b64decode(envelope['encrypted_data'])
        nonce = base64.b64decode(envelope['nonce'])
        
        # Unwrap the data key using KMS
        decrypt_response = self.kms.decrypt(
            CiphertextBlob=encrypted_dek,
            KeyId=self.kms_key_id
        )
        plaintext_dek = decrypt_response['Plaintext']
        
        # Decrypt data locally
        aesgcm = AESGCM(plaintext_dek)
        plaintext = aesgcm.decrypt(nonce, ciphertext, None)
        
        # Clean up
        plaintext_dek = b'\x00' * len(plaintext_dek)
        del plaintext_dek
        
        return plaintext
 
 
# Usage example
if __name__ == "__main__":
    # KMS key ID (ARN or alias)
    kms_key = "alias/my-master-key"
    
    env_encrypt = EnvelopeEncryption(kms_key)
    
    # Encrypt sensitive data
    secret_data = b"Customer SSN: 123-45-6789"
    envelope = env_encrypt.encrypt_data(secret_data)
    
    print(f"Encrypted envelope: {envelope}")
    
    # Later, decrypt when needed
    recovered = env_encrypt.decrypt_data(envelope)
    assert recovered == secret_data

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

import os
import ctypes
import secrets
from typing import Optional
 
class SecureKey:
    """
    A wrapper for cryptographic keys that attempts to provide
    secure memory handling. Note: Python's memory model makes
    perfect security impossible; use HSMs for critical keys.
    """
    
    def __init__(self, key_bytes: bytes):
        # Store key in a mutable bytearray so we can zero it
        self._key = bytearray(key_bytes)
        self._destroyed = False
        
        # Attempt to lock memory (platform-specific)
        self._try_lock_memory()
    
    def _try_lock_memory(self):
        """Try to prevent key from being swapped to disk."""
        try:
            if hasattr(ctypes, 'pythonapi'):
                # On Linux, we could use mlock via ctypes
                # This is illustrative - real implementation needs more care
                pass
        except Exception:
            # Fail silently - not all platforms support this
            pass
    
    def get_key(self) -> bytes:
        """Retrieve the key for cryptographic operations."""
        if self._destroyed:
            raise ValueError("Key has been destroyed")
        return bytes(self._key)
    
    def destroy(self):
        """
        Securely destroy the key by overwriting memory.
        Call this as soon as the key is no longer needed.
        """
        if not self._destroyed:
            # Overwrite with random data, then zeros
            for i in range(len(self._key)):
                self._key[i] = secrets.randbelow(256)
            for i in range(len(self._key)):
                self._key[i] = 0
            self._destroyed = True
    
    def __del__(self):
        """Destructor ensures key is wiped if object is garbage collected."""
        if hasattr(self, '_destroyed') and not self._destroyed:
            self.destroy()
    
    def __enter__(self):
        return self
    
    def __exit__(self, *args):
        self.destroy()
 
 
# Usage with context manager ensures cleanup
def encrypt_with_secure_key(data: bytes, key_bytes: bytes) -> bytes:
    """Example of using SecureKey with guaranteed cleanup."""
    with SecureKey(key_bytes) as secure_key:
        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
        
        key = secure_key.get_key()
        aesgcm = AESGCM(key)
        nonce = os.urandom(12)
        ciphertext = aesgcm.encrypt(nonce, data, None)
        
        return nonce + ciphertext
    # Key is automatically destroyed when exiting context
 
 
# Secure configuration loading pattern
def load_key_from_secure_source() -> bytes:
    """
    Load key from appropriate source based on environment.
    Never hardcode keys or use environment variables in production.
    """
    environment = os.environ.get('ENV', 'development')
    
    if environment == 'production':
        # Use AWS KMS or similar
        import boto3
        kms = boto3.client('kms')
        response = kms.generate_data_key(
            KeyId='alias/my-master-key',
            KeySpec='AES_256'
        )
        return response['Plaintext']
    
    elif environment == 'development':
        # Generate ephemeral key for development
        # NEVER use predictable or hardcoded keys
        return secrets.token_bytes(32)
    
    else:
        raise ValueError(f"Unknown environment: {environment}")

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

import boto3
 
class SecureDataService:
    """
    Demonstrates encryption context for purpose binding.
    The same key cannot decrypt data from a different context.
    """
    
    def __init__(self, kms_key_id: str):
        self.kms = boto3.client('kms')
        self.kms_key_id = kms_key_id
    
    def encrypt_user_data(self, user_id: str, data: bytes) -> dict:
        """
        Encrypt user data with context binding.
        The encryption context ties the ciphertext to this specific user
        and data classification.
        """
        encryption_context = {
            'purpose': 'user_pii_storage',
            'user_id': user_id,
            'classification': 'confidential',
            'service': 'user-service'
        }
        
        response = self.kms.encrypt(
            KeyId=self.kms_key_id,
            Plaintext=data,
            EncryptionContext=encryption_context  # This context must match on decrypt
        )
        
        return {
            'ciphertext': response['CiphertextBlob'],
            'context': encryption_context  # Store with the ciphertext
        }
    
    def decrypt_user_data(self, encrypted_package: dict) -> bytes:
        """
        Decrypt user data - must provide matching context.
        If an attacker tries to decrypt with wrong context, KMS rejects it.
        """
        response = self.kms.decrypt(
            CiphertextBlob=encrypted_package['ciphertext'],
            EncryptionContext=encrypted_package['context'],  # Must match exactly
            KeyId=self.kms_key_id
        )
        
        return response['Plaintext']
    
    def encrypt_payment_data(self, transaction_id: str, data: bytes) -> dict:
        """
        Different encryption context for payment data.
        Even with the same key, cannot decrypt as user data.
        """
        encryption_context = {
            'purpose': 'payment_processing',
            'transaction_id': transaction_id,
            'classification': 'pci_restricted',
            'service': 'payment-service'
        }
        
        response = self.kms.encrypt(
            KeyId=self.kms_key_id,
            Plaintext=data,
            EncryptionContext=encryption_context
        )
        
        return {
            'ciphertext': response['CiphertextBlob'],
            'context': encryption_context
        }
 
# This prevents the following attack:
# 1. Attacker obtains encrypted user data ciphertext
# 2. Attacker has access to payment service which can decrypt
# 3. Attacker tries to use payment service to decrypt user data
# 4. FAILS: encryption context doesn't match

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

import boto3
from datetime import datetime, timedelta
from enum import Enum
 
class KeyState(Enum):
    ACTIVE = "active"
    ROTATION_PENDING = "rotation_pending"
    ARCHIVED = "archived"
    PENDING_DELETION = "pending_deletion"
    DESTROYED = "destroyed"
 
class KeyLifecycleManager:
    """
    Manages the complete lifecycle of KMS keys including
    archival and scheduled destruction.
    """
    
    def __init__(self, region: str = 'us-east-1'):
        self.kms = boto3.client('kms', region_name=region)
    
    def get_key_metadata(self, key_id: str) -> dict:
        """Retrieve comprehensive key metadata."""
        response = self.kms.describe_key(KeyId=key_id)
        key_metadata = response['KeyMetadata']
        
        # Get key rotation status
        try:
            rotation = self.kms.get_key_rotation_status(KeyId=key_id)
            key_metadata['RotationEnabled'] = rotation['KeyRotationEnabled']
        except:
            key_metadata['RotationEnabled'] = None
        
        return key_metadata
    
    def schedule_key_deletion(self, key_id: str, waiting_days: int = 7) -> dict:
        """
        Schedule a key for deletion with the minimum safe waiting period.
        The waiting period allows cancellation if deletion was accidental.
        """
        if waiting_days < 7 or waiting_days > 30:
            raise ValueError("Waiting period must be between 7 and 30 days")
        
        # First, verify no active grants or aliases
        aliases = self.kms.list_aliases(KeyId=key_id)
        if aliases['Aliases']:
            raise ValueError(f"Key still has active aliases: {aliases['Aliases']}")
        
        # Schedule deletion
        response = self.kms.schedule_key_deletion(
            KeyId=key_id,
            PendingWindowInDays=waiting_days
        )
        
        return {
            'key_id': response['KeyId'],
            'deletion_date': response['DeletionDate'],
            'state': KeyState.PENDING_DELETION.value
        }
    
    def cancel_key_deletion(self, key_id: str) -> dict:
        """Cancel scheduled key deletion (within waiting period)."""
        response = self.kms.cancel_key_deletion(KeyId=key_id)
        
        # Re-enable the key for use
        self.kms.enable_key(KeyId=key_id)
        
        return {
            'key_id': response['KeyId'],
            'state': KeyState.ACTIVE.value
        }
    
    def archive_key(self, key_id: str, archive_reason: str) -> dict:
        """
        Archive a key by disabling it and adding archive metadata.
        The key remains but cannot be used for new operations.
        """
        # Disable the key to prevent new cryptographic operations
        self.kms.disable_key(KeyId=key_id)
        
        # Update tags to indicate archived status
        self.kms.tag_resource(
            KeyId=key_id,
            Tags=[
                {'TagKey': 'KeyState', 'TagValue': 'archived'},
                {'TagKey': 'ArchiveDate', 'TagValue': datetime.utcnow().isoformat()},
                {'TagKey': 'ArchiveReason', 'TagValue': archive_reason},
            ]
        )
        
        # Remove aliases to prevent accidental use
        aliases = self.kms.list_aliases(KeyId=key_id)
        for alias in aliases.get('Aliases', []):
            if not alias['AliasName'].startswith('alias/aws/'):
                self.kms.delete_alias(AliasName=alias['AliasName'])
        
        return {
            'key_id': key_id,
            'state': KeyState.ARCHIVED.value,
            'note': 'Key disabled, aliases removed. Re-enable temporarily for legacy decryption.'
        }
    
    def temporary_reactivate_for_decryption(self, key_id: str, duration_minutes: int = 60):
        """
        Temporarily enable an archived key for legacy data decryption.
        Should be used sparingly and logged.
        """
        # Enable the key
        self.kms.enable_key(KeyId=key_id)
        
        return {
            'key_id': key_id,
            'reactivated_until': datetime.utcnow() + timedelta(minutes=duration_minutes),
            'warning': 'Remember to manually disable after use or implement automated disable'
        }
 
 
# Example lifecycle workflow
def key_lifecycle_example():
    manager = KeyLifecycleManager()
    
    # 1. Key is in active use for encryption/decryption
    # ... (handled by application code)
    
    # 2. Key rotation period reached - archive old key
    old_key_id = "arn:aws:kms:us-east-1:123456789012:key/old-key-uuid"
    manager.archive_key(old_key_id, "Rotation: replaced by key-v2")
    
    # 3. Months later, all data re-encrypted with new key
    # Schedule old key for deletion
    manager.schedule_key_deletion(old_key_id, waiting_days=30)
    
    # 4. If we realize we still need it:
    # manager.cancel_key_deletion(old_key_id)