algo/tests/test_cloud_init_template.py

#!/usr/bin/env python3
"""
Cloud-init template validation test.

This test validates that the cloud-init template for DigitalOcean deployments
renders correctly and produces valid YAML that cloud-init can parse.

This test helps prevent regressions like issue #14800 where YAML formatting
issues caused cloud-init to fail completely, resulting in SSH timeouts.

Usage:
    python3 tests/test_cloud_init_template.py

Or from project root:
    python3 -m pytest tests/test_cloud_init_template.py -v
"""

import sys
from pathlib import Path

import yaml

# Add project root to path for imports if needed
PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

def create_expected_cloud_init():
    """
    Create the expected cloud-init content that should be generated
    by our template after the YAML indentation fix.
    """
    return """#cloud-config
# CRITICAL: The above line MUST be exactly "#cloud-config" (no space after #)
# This is required by cloud-init's YAML parser. Adding a space breaks parsing
# and causes all cloud-init directives to be skipped, resulting in SSH timeouts.
# See: https://github.com/trailofbits/algo/issues/14800
output: {all: '| tee -a /var/log/cloud-init-output.log'}

package_update: true
package_upgrade: true

packages:
  - sudo

users:
  - default
  - name: algo
    homedir: /home/algo
    sudo: ALL=(ALL) NOPASSWD:ALL
    groups: adm,netdev
    shell: /bin/bash
    lock_passwd: true
    ssh_authorized_keys:
      - "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAACAQDTest algo-test"

write_files:
  - path: /etc/ssh/sshd_config
    content: |
      Port 4160
      AllowGroups algo
      PermitRootLogin no
      PasswordAuthentication no
      ChallengeResponseAuthentication no
      UsePAM yes
      X11Forwarding yes
      PrintMotd no
      AcceptEnv LANG LC_*
      Subsystem	sftp	/usr/lib/openssh/sftp-server

runcmd:
  - set -x
  - ufw --force reset
  - sudo apt-get remove -y --purge sshguard || true
  - systemctl restart sshd.service
"""

class TestCloudInitTemplate:
    """Test class for cloud-init template validation."""

    def test_yaml_validity(self):
        """Test that the expected cloud-init YAML is valid."""
        print("🧪 Testing YAML validity...")

        cloud_init_content = create_expected_cloud_init()

        try:
            parsed = yaml.safe_load(cloud_init_content)
            print("✅ YAML parsing successful")
            assert parsed is not None, "YAML should parse to a non-None value"
            return parsed
        except yaml.YAMLError as e:
            print(f"❌ YAML parsing failed: {e}")
            assert False, f"YAML parsing failed: {e}"

    def test_required_sections(self):
        """Test that all required cloud-init sections are present."""
        print("🧪 Testing required sections...")

        parsed = self.test_yaml_validity()

        required_sections = [
            'package_update', 'package_upgrade', 'packages',
            'users', 'write_files', 'runcmd'
        ]

        missing = [section for section in required_sections if section not in parsed]
        assert not missing, f"Missing required sections: {missing}"

        print("✅ All required sections present")

    def test_ssh_configuration(self):
        """Test that SSH configuration is correct."""
        print("🧪 Testing SSH configuration...")

        parsed = self.test_yaml_validity()

        write_files = parsed.get('write_files', [])
        assert write_files, "write_files section should be present"

        # Find sshd_config file
        sshd_config = None
        for file_entry in write_files:
            if file_entry.get('path') == '/etc/ssh/sshd_config':
                sshd_config = file_entry
                break

        assert sshd_config, "sshd_config file should be in write_files"

        content = sshd_config.get('content', '')
        assert content, "sshd_config should have content"

        # Check required SSH configurations
        required_configs = [
            'Port 4160',
            'AllowGroups algo',
            'PermitRootLogin no',
            'PasswordAuthentication no'
        ]

        missing = [config for config in required_configs if config not in content]
        assert not missing, f"Missing SSH configurations: {missing}"

        # Verify proper formatting - first line should be Port directive
        lines = content.strip().split('\n')
        assert lines[0].strip() == 'Port 4160', f"First line should be 'Port 4160', got: {repr(lines[0])}"

        print("✅ SSH configuration correct")

    def test_user_creation(self):
        """Test that algo user will be created correctly."""
        print("🧪 Testing user creation...")

        parsed = self.test_yaml_validity()

        users = parsed.get('users', [])
        assert users, "users section should be present"

        # Find algo user
        algo_user = None
        for user in users:
            if isinstance(user, dict) and user.get('name') == 'algo':
                algo_user = user
                break

        assert algo_user, "algo user should be defined"

        # Check required user properties
        required_props = ['sudo', 'groups', 'shell', 'ssh_authorized_keys']
        missing = [prop for prop in required_props if prop not in algo_user]
        assert not missing, f"algo user missing properties: {missing}"

        # Verify sudo configuration
        sudo_config = algo_user.get('sudo', '')
        assert 'NOPASSWD:ALL' in sudo_config, f"sudo config should allow passwordless access: {sudo_config}"

        print("✅ User creation correct")

    def test_runcmd_section(self):
        """Test that runcmd section will restart SSH correctly."""
        print("🧪 Testing runcmd section...")

        parsed = self.test_yaml_validity()

        runcmd = parsed.get('runcmd', [])
        assert runcmd, "runcmd section should be present"

        # Check for SSH restart command
        ssh_restart_found = False
        for cmd in runcmd:
            if 'systemctl restart sshd' in str(cmd):
                ssh_restart_found = True
                break

        assert ssh_restart_found, f"SSH restart command not found in runcmd: {runcmd}"

        print("✅ runcmd section correct")

    def test_indentation_consistency(self):
        """Test that sshd_config content has consistent indentation."""
        print("🧪 Testing indentation consistency...")

        cloud_init_content = create_expected_cloud_init()

        # Extract the sshd_config content lines
        lines = cloud_init_content.split('\n')
        in_sshd_content = False
        sshd_lines = []

        for line in lines:
            if 'content: |' in line:
                in_sshd_content = True
                continue
            elif in_sshd_content:
                if line.strip() == '' and len(sshd_lines) > 0:
                    break
                if line.startswith('runcmd:'):
                    break
                sshd_lines.append(line)

        assert sshd_lines, "Should be able to extract sshd_config content"

        # Check that all non-empty lines have consistent 6-space indentation
        non_empty_lines = [line for line in sshd_lines if line.strip()]
        assert non_empty_lines, "sshd_config should have content"

        for line in non_empty_lines:
            # Each line should start with exactly 6 spaces
            assert line.startswith('      ') and not line.startswith('       '), \
                f"Line should have exactly 6 spaces indentation: {repr(line)}"

        print("✅ Indentation is consistent")

def run_tests():
    """Run all tests manually (for non-pytest usage)."""
    print("🚀 Cloud-init template validation tests")
    print("=" * 50)

    test_instance = TestCloudInitTemplate()

    try:
        test_instance.test_yaml_validity()
        test_instance.test_required_sections()
        test_instance.test_ssh_configuration()
        test_instance.test_user_creation()
        test_instance.test_runcmd_section()
        test_instance.test_indentation_consistency()

        print("=" * 50)
        print("🎉 ALL TESTS PASSED!")
        print("✅ Cloud-init template is working correctly")
        print("✅ DigitalOcean deployment should succeed")
        return True

    except AssertionError as e:
        print(f"❌ Test failed: {e}")
        return False
    except Exception as e:
        print(f"❌ Unexpected error: {e}")
        return False

if __name__ == "__main__":
    success = run_tests()
    sys.exit(0 if success else 1)