chorus-services/modules/shhh/automation/revocation.py

"""
Automated Secret Revocation System for SHHH Secrets Sentinel
Provides automated response capabilities for different types of detected secrets.
"""

import asyncio
import aiohttp
import json
import time
from typing import Dict, Any, Optional, List
from dataclasses import dataclass
from datetime import datetime
import structlog

from ..core.quarantine import QuarantineEntry, RevocationEvent, QuarantineManager
from ..core.detector import SecretMatch

logger = structlog.get_logger()


@dataclass
class RevocationRequest:
    """Represents a request to revoke a secret"""
    quarantine_id: int
    secret_type: str
    redacted_secret: str
    urgency: str  # 'immediate', 'high', 'medium', 'low'
    metadata: Dict[str, Any]


@dataclass
class RevocationResponse:
    """Represents the response from a revocation attempt"""
    success: bool
    method: str
    response_data: Dict[str, Any]
    error_message: Optional[str] = None
    revocation_id: Optional[str] = None


class SecretRevoker:
    """
    Automated secret revocation system that integrates with various cloud providers
    and services to automatically disable compromised credentials.
    """
    
    def __init__(self, quarantine_manager: QuarantineManager, webhook_config: Dict[str, str] = None):
        self.quarantine = quarantine_manager
        self.webhook_config = webhook_config or {}
        
        # Revocation timeouts and retry settings
        self.request_timeout = 10  # seconds
        self.max_retries = 3
        self.retry_delay = 2  # seconds
        
        # Statistics
        self.stats = {
            'total_revocations': 0,
            'successful_revocations': 0,
            'failed_revocations': 0,
            'revocations_by_type': {},
            'last_reset': datetime.now()
        }
        
        logger.info("Initialized SecretRevoker")
    
    async def trigger_revocation(self, quarantine_entry: QuarantineEntry) -> Optional[RevocationResponse]:
        """Trigger automatic revocation for a quarantined secret"""
        try:
            revocation_request = RevocationRequest(
                quarantine_id=quarantine_entry.id,
                secret_type=quarantine_entry.secret_type,
                redacted_secret=self._extract_redacted_from_metadata(quarantine_entry),
                urgency=self._determine_urgency(quarantine_entry.severity),
                metadata={
                    'source_agent': quarantine_entry.source_agent,
                    'detection_timestamp': quarantine_entry.timestamp.isoformat(),
                    'confidence': quarantine_entry.confidence,
                    'hypercore_position': quarantine_entry.hypercore_position
                }
            )
            
            # Determine revocation method
            revocation_method = self._get_revocation_method(quarantine_entry.secret_type)
            if not revocation_method:
                logger.warning(f"No revocation method configured for {quarantine_entry.secret_type}")
                return None
            
            # Attempt revocation
            response = await self._execute_revocation(revocation_request, revocation_method)
            
            # Record the revocation event
            await self._record_revocation_event(quarantine_entry, response)
            
            # Update statistics
            self._update_stats(quarantine_entry.secret_type, response.success)
            
            return response
            
        except Exception as e:
            logger.error(f"Failed to trigger revocation for quarantine {quarantine_entry.id}: {e}")
            return None
    
    def _extract_redacted_from_metadata(self, quarantine_entry: QuarantineEntry) -> str:
        """Extract redacted secret from quarantine metadata"""
        try:
            matches = quarantine_entry.metadata.get('matches', [])
            if matches:
                # Get the first match's redacted text
                return matches[0].get('redacted_text', 'REDACTED')
        except:
            pass
        
        return 'REDACTED'
    
    def _determine_urgency(self, severity: str) -> str:
        """Determine revocation urgency based on severity"""
        urgency_map = {
            'CRITICAL': 'immediate',
            'HIGH': 'high',
            'MEDIUM': 'medium',
            'LOW': 'low'
        }
        return urgency_map.get(severity, 'medium')
    
    def _get_revocation_method(self, secret_type: str) -> Optional[str]:
        """Get the appropriate revocation method for a secret type"""
        method_map = {
            'AWS_ACCESS_KEY': 'aws_iam_revocation',
            'AWS_SECRET_KEY': 'aws_iam_revocation',
            'GITHUB_TOKEN': 'github_token_revocation',
            'GITHUB_OAUTH': 'github_token_revocation',
            'SLACK_TOKEN': 'slack_token_revocation',
            'GOOGLE_API_KEY': 'google_api_revocation',
            'DOCKER_TOKEN': 'docker_token_revocation'
        }
        return method_map.get(secret_type)
    
    async def _execute_revocation(self, request: RevocationRequest, method: str) -> RevocationResponse:
        """Execute the actual revocation based on the method"""
        method_handlers = {
            'aws_iam_revocation': self._revoke_aws_credentials,
            'github_token_revocation': self._revoke_github_token,
            'slack_token_revocation': self._revoke_slack_token,
            'google_api_revocation': self._revoke_google_api_key,
            'docker_token_revocation': self._revoke_docker_token,
            'webhook_revocation': self._revoke_via_webhook
        }
        
        handler = method_handlers.get(method, self._revoke_via_webhook)
        
        for attempt in range(self.max_retries):
            try:
                response = await handler(request)
                if response.success:
                    logger.info(
                        f"Successfully revoked {request.secret_type}",
                        quarantine_id=request.quarantine_id,
                        method=method,
                        attempt=attempt + 1
                    )
                    return response
                
                # Log failure and retry if not successful
                logger.warning(
                    f"Revocation attempt {attempt + 1} failed",
                    quarantine_id=request.quarantine_id,
                    method=method,
                    error=response.error_message
                )
                
                if attempt < self.max_retries - 1:
                    await asyncio.sleep(self.retry_delay * (attempt + 1))  # Exponential backoff
                
            except Exception as e:
                logger.error(f"Revocation attempt {attempt + 1} error: {e}")
                if attempt < self.max_retries - 1:
                    await asyncio.sleep(self.retry_delay * (attempt + 1))
        
        # All attempts failed
        return RevocationResponse(
            success=False,
            method=method,
            response_data={},
            error_message=f"All {self.max_retries} revocation attempts failed"
        )
    
    async def _revoke_aws_credentials(self, request: RevocationRequest) -> RevocationResponse:
        """Revoke AWS credentials via webhook"""
        webhook_url = self.webhook_config.get('AWS_ACCESS_KEY')
        if not webhook_url:
            return RevocationResponse(
                success=False,
                method='aws_iam_revocation',
                response_data={},
                error_message="No AWS revocation webhook configured"
            )
        
        payload = {
            'event': 'secret_leak_detected',
            'secret_type': request.secret_type,
            'redacted_key': request.redacted_secret,
            'urgency': request.urgency,
            'quarantine_id': request.quarantine_id,
            'timestamp': datetime.now().isoformat(),
            'recommended_action': 'Revoke IAM access key immediately',
            'metadata': request.metadata
        }
        
        return await self._send_webhook_request(webhook_url, payload, 'aws_iam_revocation')
    
    async def _revoke_github_token(self, request: RevocationRequest) -> RevocationResponse:
        """Revoke GitHub token via webhook"""
        webhook_url = self.webhook_config.get('GITHUB_TOKEN')
        if not webhook_url:
            return RevocationResponse(
                success=False,
                method='github_token_revocation',
                response_data={},
                error_message="No GitHub revocation webhook configured"
            )
        
        payload = {
            'event': 'secret_leak_detected',
            'secret_type': request.secret_type,
            'redacted_key': request.redacted_secret,
            'urgency': request.urgency,
            'quarantine_id': request.quarantine_id,
            'timestamp': datetime.now().isoformat(),
            'recommended_action': 'Revoke GitHub token via API or settings',
            'metadata': request.metadata
        }
        
        return await self._send_webhook_request(webhook_url, payload, 'github_token_revocation')
    
    async def _revoke_slack_token(self, request: RevocationRequest) -> RevocationResponse:
        """Revoke Slack token via webhook"""
        webhook_url = self.webhook_config.get('SLACK_TOKEN')
        if not webhook_url:
            return RevocationResponse(
                success=False,
                method='slack_token_revocation',
                response_data={},
                error_message="No Slack revocation webhook configured"
            )
        
        payload = {
            'event': 'secret_leak_detected',
            'secret_type': request.secret_type,
            'redacted_key': request.redacted_secret,
            'urgency': request.urgency,
            'quarantine_id': request.quarantine_id,
            'timestamp': datetime.now().isoformat(),
            'recommended_action': 'Revoke Slack token via Admin API',
            'metadata': request.metadata
        }
        
        return await self._send_webhook_request(webhook_url, payload, 'slack_token_revocation')
    
    async def _revoke_google_api_key(self, request: RevocationRequest) -> RevocationResponse:
        """Revoke Google API key via webhook"""
        webhook_url = self.webhook_config.get('GOOGLE_API_KEY')
        if not webhook_url:
            return RevocationResponse(
                success=False,
                method='google_api_revocation',
                response_data={},
                error_message="No Google API revocation webhook configured"
            )
        
        payload = {
            'event': 'secret_leak_detected',
            'secret_type': request.secret_type,
            'redacted_key': request.redacted_secret,
            'urgency': request.urgency,
            'quarantine_id': request.quarantine_id,
            'timestamp': datetime.now().isoformat(),
            'recommended_action': 'Revoke API key via Google Cloud Console',
            'metadata': request.metadata
        }
        
        return await self._send_webhook_request(webhook_url, payload, 'google_api_revocation')
    
    async def _revoke_docker_token(self, request: RevocationRequest) -> RevocationResponse:
        """Revoke Docker token via webhook"""
        webhook_url = self.webhook_config.get('DOCKER_TOKEN')
        if not webhook_url:
            return RevocationResponse(
                success=False,
                method='docker_token_revocation',
                response_data={},
                error_message="No Docker revocation webhook configured"
            )
        
        payload = {
            'event': 'secret_leak_detected',
            'secret_type': request.secret_type,
            'redacted_key': request.redacted_secret,
            'urgency': request.urgency,
            'quarantine_id': request.quarantine_id,
            'timestamp': datetime.now().isoformat(),
            'recommended_action': 'Revoke Docker token via Hub settings',
            'metadata': request.metadata
        }
        
        return await self._send_webhook_request(webhook_url, payload, 'docker_token_revocation')
    
    async def _revoke_via_webhook(self, request: RevocationRequest) -> RevocationResponse:
        """Generic webhook revocation for unknown secret types"""
        # Try to find a generic webhook endpoint
        webhook_url = self.webhook_config.get('GENERIC', 
                                               self.webhook_config.get('DEFAULT'))
        
        if not webhook_url:
            return RevocationResponse(
                success=False,
                method='webhook_revocation',
                response_data={},
                error_message=f"No webhook configured for {request.secret_type}"
            )
        
        payload = {
            'event': 'secret_leak_detected',
            'secret_type': request.secret_type,
            'redacted_key': request.redacted_secret,
            'urgency': request.urgency,
            'quarantine_id': request.quarantine_id,
            'timestamp': datetime.now().isoformat(),
            'recommended_action': 'Manual review and revocation required',
            'metadata': request.metadata
        }
        
        return await self._send_webhook_request(webhook_url, payload, 'webhook_revocation')
    
    async def _send_webhook_request(self, url: str, payload: Dict[str, Any], method: str) -> RevocationResponse:
        """Send webhook request and handle response"""
        try:
            async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=self.request_timeout)) as session:
                async with session.post(url, json=payload) as response:
                    response_data = {}
                    try:
                        response_data = await response.json()
                    except:
                        response_data = {'text': await response.text()}
                    
                    if response.status == 200:
                        return RevocationResponse(
                            success=True,
                            method=method,
                            response_data=response_data,
                            revocation_id=response_data.get('revocation_id')
                        )
                    else:
                        return RevocationResponse(
                            success=False,
                            method=method,
                            response_data=response_data,
                            error_message=f"HTTP {response.status}: {response_data}"
                        )
                        
        except asyncio.TimeoutError:
            return RevocationResponse(
                success=False,
                method=method,
                response_data={},
                error_message=f"Webhook request timed out after {self.request_timeout}s"
            )
        except Exception as e:
            return RevocationResponse(
                success=False,
                method=method,
                response_data={},
                error_message=f"Webhook request failed: {str(e)}"
            )
    
    async def _record_revocation_event(self, quarantine_entry: QuarantineEntry, response: RevocationResponse):
        """Record revocation event in the database"""
        try:
            revocation_event = RevocationEvent(
                quarantine_id=quarantine_entry.id,
                secret_type=quarantine_entry.secret_type,
                revocation_method=response.method,
                status='success' if response.success else 'failed',
                response_data=response.response_data,
                timestamp=datetime.now()
            )
            
            await self.quarantine.record_revocation(revocation_event)
            
        except Exception as e:
            logger.error(f"Failed to record revocation event: {e}")
    
    def _update_stats(self, secret_type: str, success: bool):
        """Update revocation statistics"""
        self.stats['total_revocations'] += 1
        
        if success:
            self.stats['successful_revocations'] += 1
        else:
            self.stats['failed_revocations'] += 1
        
        # Update by-type stats
        if secret_type not in self.stats['revocations_by_type']:
            self.stats['revocations_by_type'][secret_type] = {
                'total': 0,
                'successful': 0,
                'failed': 0
            }
        
        type_stats = self.stats['revocations_by_type'][secret_type]
        type_stats['total'] += 1
        
        if success:
            type_stats['successful'] += 1
        else:
            type_stats['failed'] += 1
    
    async def test_webhook_endpoint(self, secret_type: str) -> Dict[str, Any]:
        """Test a webhook endpoint with a test payload"""
        webhook_url = self.webhook_config.get(secret_type)
        if not webhook_url:
            return {
                'success': False,
                'error': f'No webhook configured for {secret_type}'
            }
        
        test_payload = {
            'event': 'webhook_test',
            'secret_type': secret_type,
            'test': True,
            'timestamp': datetime.now().isoformat()
        }
        
        try:
            response = await self._send_webhook_request(webhook_url, test_payload, 'test')
            return {
                'success': response.success,
                'method': response.method,
                'response_data': response.response_data,
                'error': response.error_message
            }
        except Exception as e:
            return {
                'success': False,
                'error': str(e)
            }
    
    def get_stats(self) -> Dict[str, Any]:
        """Get revocation statistics"""
        current_time = datetime.now()
        uptime_hours = (current_time - self.stats['last_reset']).total_seconds() / 3600
        
        stats = self.stats.copy()
        stats.update({
            'uptime_hours': round(uptime_hours, 2),
            'success_rate': (
                self.stats['successful_revocations'] / max(1, self.stats['total_revocations'])
            ) * 100,
            'configured_webhooks': list(self.webhook_config.keys())
        })
        
        return stats
    
    def reset_stats(self):
        """Reset statistics counters"""
        self.stats = {
            'total_revocations': 0,
            'successful_revocations': 0,
            'failed_revocations': 0,
            'revocations_by_type': {},
            'last_reset': datetime.now()
        }
        
        logger.info("SecretRevoker statistics reset")