ROLES = {
    'DataEngineer': [
        (permissions.ACTION_CAN_READ, permissions.RESOURCE_DAG),
        (permissions.ACTION_CAN_EDIT, permissions.RESOURCE_DAG),
    ],
    'DataAnalyst': [
        (permissions.ACTION_CAN_READ, permissions.RESOURCE_DAG),
    ],
}

# airflow.cfg
[secrets]
backend = airflow.providers.hashicorp.secrets.vault.VaultBackend
backend_kwargs = {
    "connections_path": "airflow/connections",
    "variables_path": "airflow/variables",
}

# Encrypted connection example
conn = Connection(
    conn_id='production_db',
    conn_type='postgres',
    host='db.example.com',
    extra='{"sslmode": "require", "sslrootcert": "/path/to/ca.pem"}',
)

# Additional SSL parameters for strict verification
conn_params = {
    'sslmode': 'require',  # Enforce SSL
    'sslcert': '/path/to/client-cert.pem',
    'sslkey': '/path/to/client-key.pem',
    'sslrootcert': '/path/to/ca-cert.pem',
}

def verify_encryption():
    """Verify connection is encrypted."""
    hook = PostgresHook(postgres_conn_id='encrypted_db')

    result = hook.get_first("""
        SELECT
            pgssl.ssl_is_used as ssl_enabled,
            pgssl.ssl_version as ssl_version,
            pgssl.ssl_cipher as cipher
        FROM pg_stat_ssl pgssl
        JOIN pg_stat_activity pgact ON pgssl.pid = pgact.pid
        WHERE pgact.usename = current_user
    """)

    return {
        'ssl_enabled': result[0],
        'ssl_version': result[1],
        'cipher': result[2],
    }

# security_audit.py
from airflow import settings
from airflow.models import Connection, Variable
from sqlalchemy import text
import json

def audit_security_posture():
    """Comprehensive security audit of Airflow deployment."""
    session = settings.Session()
    findings = []
    
    # Check for plaintext passwords in metadata DB
    plaintext_conns = session.query(Connection).filter(
        Connection.conn_type.notin_(['aws', 'google_cloud_platform'])
    ).all()
    
    for conn in plaintext_conns:
        if conn.password and not conn.password.startswith('{'):
            findings.append({
                'severity': 'HIGH',
                'category': 'Secrets',
                'finding': f'Plaintext password in connection: {conn.conn_id}',
                'recommendation': 'Move to secrets backend',
            })
    
    # Check for variables with sensitive data
    sensitive_patterns = ['password', 'secret', 'key', 'token']
    variables = session.query(Variable).all()
    
    for var in variables:
        if any(pattern in var.key.lower() for pattern in sensitive_patterns):
            findings.append({
                'severity': 'MEDIUM',
                'category': 'Secrets',
                'finding': f'Potentially sensitive variable: {var.key}',
                'recommendation': 'Move to secrets backend',
            })
    
    # Check for DAGs with hardcoded secrets
    import ast
    import os
    
    dags_folder = '/opt/airflow/dags'
    for root, dirs, files in os.walk(dags_folder):
        for file in files:
            if file.endswith('.py'):
                filepath = os.path.join(root, file)
                with open(filepath, 'r') as f:
                    content = f.read()
                
                # Simple pattern matching
                for pattern in ['password=', 'secret=', 'api_key=']:
                    if pattern in content:
                        findings.append({
                            'severity': 'HIGH',
                            'category': 'Code',
                            'finding': f'Potential hardcoded secret in {filepath}',
                            'recommendation': 'Use variables or secrets backend',
                        })
    
    return findings

def generate_security_report(findings):
    """Generate security audit report."""
    report = {
        'total_findings': len(findings),
        'high_severity': len([f for f in findings if f['severity'] == 'HIGH']),
        'medium_severity': len([f for f in findings if f['severity'] == 'MEDIUM']),
        'low_severity': len([f for f in findings if f['severity'] == 'LOW']),
        'findings': findings,
    }
    
    print(json.dumps(report, indent=2))
    return report

if __name__ == "__main__":
    findings = audit_security_posture()
    generate_security_report(findings)