""" 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)}")