headplane/remote-access/token_encryption.py

"""
Heady Remote Access - Token Encryption
Secure AES-256-CBC encryption for guacamole-lite connection tokens
Based on guacamole-lite encryption pattern for maximum security
"""
import os
import json
import base64
import secrets
from typing import Dict, Any
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import padding
import logging

logger = logging.getLogger(__name__)

class GuacamoleTokenEncryption:
    """
    Secure token encryption for Guacamole connections
    Uses AES-256-CBC with random IVs for maximum security
    """
    
    def __init__(self, secret_key: str = None):
        """
        Initialize encryption with secret key
        Secret key should be 32 bytes (256 bits) for AES-256
        """
        if secret_key is None:
            secret_key = os.getenv("GUACAMOLE_ENCRYPTION_KEY")
            
        if not secret_key:
            raise ValueError("GUACAMOLE_ENCRYPTION_KEY environment variable required")
            
        # Ensure key is exactly 32 bytes for AES-256
        if len(secret_key.encode()) != 32:
            # Hash or pad the key to exactly 32 bytes
            import hashlib
            self.secret_key = hashlib.sha256(secret_key.encode()).digest()
        else:
            self.secret_key = secret_key.encode()
    
    def encrypt_token(self, connection_params: Dict[str, Any]) -> str:
        """
        Encrypt connection parameters into a secure token
        
        Args:
            connection_params: Dictionary containing Guacamole connection settings
            
        Returns:
            Base64-encoded encrypted token
        """
        try:
            # Convert params to JSON string
            params_json = json.dumps(connection_params, separators=(',', ':'))
            params_bytes = params_json.encode('utf-8')
            
            # Generate random IV (16 bytes for AES-CBC)
            iv = secrets.token_bytes(16)
            
            # Create cipher
            cipher = Cipher(
                algorithms.AES(self.secret_key),
                modes.CBC(iv),
                backend=default_backend()
            )
            encryptor = cipher.encryptor()
            
            # Pad data to block size (16 bytes for AES)
            padder = padding.PKCS7(128).padder()  # 128 bits = 16 bytes
            padded_data = padder.update(params_bytes)
            padded_data += padder.finalize()
            
            # Encrypt the data
            encrypted_data = encryptor.update(padded_data) + encryptor.finalize()
            
            # Combine IV + encrypted data and encode as base64
            token_data = iv + encrypted_data
            encrypted_token = base64.b64encode(token_data).decode('utf-8')
            
            logger.info(f"Generated encrypted token for {connection_params.get('protocol', 'unknown')} connection")
            return encrypted_token
            
        except Exception as e:
            logger.error(f"Token encryption failed: {e}")
            raise
    
    def decrypt_token(self, encrypted_token: str) -> Dict[str, Any]:
        """
        Decrypt a token back to connection parameters
        (For testing/debugging purposes only)
        
        Args:
            encrypted_token: Base64-encoded encrypted token
            
        Returns:
            Decrypted connection parameters
        """
        try:
            # Decode base64
            token_data = base64.b64decode(encrypted_token.encode('utf-8'))
            
            # Extract IV (first 16 bytes) and encrypted data
            iv = token_data[:16]
            encrypted_data = token_data[16:]
            
            # Create cipher
            cipher = Cipher(
                algorithms.AES(self.secret_key),
                modes.CBC(iv),
                backend=default_backend()
            )
            decryptor = cipher.decryptor()
            
            # Decrypt
            padded_data = decryptor.update(encrypted_data) + decryptor.finalize()
            
            # Remove padding
            unpadder = padding.PKCS7(128).unpadder()
            params_bytes = unpadder.update(padded_data)
            params_bytes += unpadder.finalize()
            
            # Parse JSON
            params_json = params_bytes.decode('utf-8')
            connection_params = json.loads(params_json)
            
            return connection_params
            
        except Exception as e:
            logger.error(f"Token decryption failed: {e}")
            raise ValueError("Invalid or corrupted token")


def create_connection_token(
    protocol: str,
    hostname: str,
    username: str,
    port: int = None,
    additional_params: Dict[str, Any] = None
) -> str:
    """
    Create encrypted connection token for Heady remote access
    
    Args:
        protocol: Connection protocol (ssh, rdp, vnc, telnet, kubernetes)
        hostname: Target hostname/IP
        username: Username for connection
        port: Port number (optional, protocol defaults used)
        additional_params: Extra protocol-specific parameters
        
    Returns:
        Encrypted token string for guacamole-lite
    """
    
    # Base connection parameters
    connection_params = {
        "connection": {
            "type": protocol,
            "settings": {
                "hostname": hostname,
                "username": username,
                "port": str(port or get_default_port(protocol))
            }
        }
    }
    
    # Add protocol-specific settings
    if protocol == "ssh":
        connection_params["connection"]["settings"].update({
            "font-size": "12",
            "color-scheme": "gray-black",
            "enable-sftp": "true",
            "sftp-root-directory": f"/home/{username}"
        })
        
    elif protocol == "rdp":
        connection_params["connection"]["settings"].update({
            "security": "any",
            "ignore-cert": "true", 
            "enable-drive": "true",
            "drive-path": "/drive",
            "console": "true",
            "width": "1920",
            "height": "1080",
            "color-depth": "24"
        })
        
    elif protocol == "vnc":
        connection_params["connection"]["settings"].update({
            "color-depth": "24",
            "cursor": "remote"
        })
        
    elif protocol == "telnet":
        connection_params["connection"]["settings"].update({
            "font-size": "12",
            "color-scheme": "green-black",
            "backspace": "127"
        })
        
    elif protocol == "kubernetes":
        # Special case for kubectl exec
        connection_params["connection"]["settings"].update({
            "font-size": "12",
            "color-scheme": "blue-black",
            "command": f"kubectl exec -it {hostname} -- /bin/bash"
        })
    
    # Add any additional parameters
    if additional_params:
        connection_params["connection"]["settings"].update(additional_params)
    
    # Add recording settings if needed
    recording_path = additional_params.get("recording_path") if additional_params else None
    if recording_path:
        connection_params["connection"]["settings"].update({
            "recording-path": recording_path,
            "recording-name": f"session-{secrets.token_hex(8)}",
            "create-recording-path": "true"
        })
    
    # Encrypt the token
    encryptor = GuacamoleTokenEncryption()
    return encryptor.encrypt_token(connection_params)


def get_default_port(protocol: str) -> int:
    """Get default port for protocol"""
    ports = {
        "ssh": 22,
        "rdp": 3389, 
        "vnc": 5901,
        "telnet": 23,
        "kubernetes": 22  # Usually SSH to k8s node
    }
    return ports.get(protocol, 22)


def validate_connection_token(token: str) -> bool:
    """
    Validate that a token can be decrypted successfully
    (For testing/validation purposes)
    """
    try:
        encryptor = GuacamoleTokenEncryption()
        params = encryptor.decrypt_token(token)
        return "connection" in params and "settings" in params["connection"]
    except:
        return False


# Example usage for testing
if __name__ == "__main__":
    # Set test encryption key
    os.environ["GUACAMOLE_ENCRYPTION_KEY"] = "MySuperSecretKeyForHeadyTokens32"
    
    # Test SSH connection token
    ssh_token = create_connection_token(
        protocol="ssh",
        hostname="test-node.tailscale.net",
        username="heady-user",
        additional_params={
            "recording_path": "/recordings",
            "enable-sftp": "true"
        }
    )
    
    print(f"SSH Token: {ssh_token}")
    print(f"Token valid: {validate_connection_token(ssh_token)}")
    
    # Test RDP connection token
    rdp_token = create_connection_token(
        protocol="rdp",
        hostname="windows-box.tailscale.net", 
        username="Administrator",
        additional_params={
            "recording_path": "/recordings",
            "width": "1920",
            "height": "1080"
        }
    )
    
    print(f"RDP Token: {rdp_token}")
    print(f"Token valid: {validate_connection_token(rdp_token)}")