auto_sigma_rule_generator/backend/tasks/bulk_tasks.py

"""
Bulk processing tasks for Celery
"""
import asyncio
import logging
from typing import Optional, Dict, Any
from celery import current_task
from celery_config import celery_app, get_db_session
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from bulk_seeder import BulkSeeder

logger = logging.getLogger(__name__)

@celery_app.task(bind=True, name='bulk_tasks.full_bulk_seed')
def full_bulk_seed_task(self, start_year: int = 2002, end_year: Optional[int] = None, 
                       skip_nvd: bool = False, skip_nomi_sec: bool = False, 
                       skip_exploitdb: bool = False, skip_cisa_kev: bool = False) -> Dict[str, Any]:
    """
    Celery task for full bulk seeding operation
    
    Args:
        start_year: Starting year for NVD data
        end_year: Ending year for NVD data
        skip_nvd: Skip NVD bulk processing
        skip_nomi_sec: Skip nomi-sec PoC synchronization
        skip_exploitdb: Skip ExploitDB synchronization
        skip_cisa_kev: Skip CISA KEV synchronization
        
    Returns:
        Dictionary containing operation results
    """
    db_session = get_db_session()
    
    try:
        # Update task progress
        self.update_state(
            state='PROGRESS', 
            meta={
                'stage': 'initializing',
                'progress': 0,
                'message': 'Starting bulk seeding operation'
            }
        )
        
        logger.info(f"Starting full bulk seed task: {start_year}-{end_year}")
        
        # Create seeder instance
        seeder = BulkSeeder(db_session)
        
        # Create progress callback
        def update_progress(stage: str, progress: int, message: str = None):
            self.update_state(
                state='PROGRESS',
                meta={
                    'stage': stage,
                    'progress': progress,
                    'message': message or f'Processing {stage}'
                }
            )
        
        # Run the bulk seeding operation
        # Note: We need to handle the async nature of bulk_seeder
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        
        try:
            result = loop.run_until_complete(
                seeder.full_bulk_seed(
                    start_year=start_year,
                    end_year=end_year,
                    skip_nvd=skip_nvd,
                    skip_nomi_sec=skip_nomi_sec,
                    skip_exploitdb=skip_exploitdb,
                    skip_cisa_kev=skip_cisa_kev,
                    progress_callback=update_progress
                )
            )
        finally:
            loop.close()
        
        # Update final progress
        self.update_state(
            state='SUCCESS',
            meta={
                'stage': 'completed',
                'progress': 100,
                'message': 'Bulk seeding completed successfully'
            }
        )
        
        logger.info(f"Full bulk seed task completed: {result}")
        return result
        
    except Exception as e:
        logger.error(f"Full bulk seed task failed: {e}")
        self.update_state(
            state='FAILURE',
            meta={
                'stage': 'error',
                'progress': 0,
                'message': f'Task failed: {str(e)}',
                'error': str(e)
            }
        )
        raise
    finally:
        db_session.close()

@celery_app.task(bind=True, name='bulk_tasks.incremental_update_task')
def incremental_update_task(self) -> Dict[str, Any]:
    """
    Celery task for incremental updates
    
    Returns:
        Dictionary containing update results
    """
    db_session = get_db_session()
    
    try:
        # Update task progress
        self.update_state(
            state='PROGRESS',
            meta={
                'stage': 'incremental_update',
                'progress': 0,
                'message': 'Starting incremental update'
            }
        )
        
        logger.info("Starting incremental update task")
        
        # Create seeder instance
        seeder = BulkSeeder(db_session)
        
        # Run the incremental update
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        
        try:
            result = loop.run_until_complete(seeder.incremental_update())
        finally:
            loop.close()
        
        # Update final progress
        self.update_state(
            state='SUCCESS',
            meta={
                'stage': 'completed',
                'progress': 100,
                'message': 'Incremental update completed successfully'
            }
        )
        
        logger.info(f"Incremental update task completed: {result}")
        return result
        
    except Exception as e:
        logger.error(f"Incremental update task failed: {e}")
        self.update_state(
            state='FAILURE',
            meta={
                'stage': 'error',
                'progress': 0,
                'message': f'Task failed: {str(e)}',
                'error': str(e)
            }
        )
        raise
    finally:
        db_session.close()

@celery_app.task(bind=True, name='bulk_tasks.generate_enhanced_sigma_rules')
def generate_enhanced_sigma_rules_task(self) -> Dict[str, Any]:
    """
    Celery task for generating enhanced SIGMA rules
    
    Returns:
        Dictionary containing generation results
    """
    db_session = get_db_session()
    
    try:
        # Update task progress
        self.update_state(
            state='PROGRESS',
            meta={
                'stage': 'generating_rules',
                'progress': 0,
                'message': 'Starting enhanced SIGMA rule generation'
            }
        )
        
        logger.info("Starting enhanced SIGMA rule generation task")
        
        # Create seeder instance
        seeder = BulkSeeder(db_session)
        
        # Run the rule generation
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        
        try:
            result = loop.run_until_complete(seeder.generate_enhanced_sigma_rules())
        finally:
            loop.close()
        
        # Update final progress
        self.update_state(
            state='SUCCESS',
            meta={
                'stage': 'completed',
                'progress': 100,
                'message': 'Enhanced SIGMA rule generation completed successfully'
            }
        )
        
        logger.info(f"Enhanced SIGMA rule generation task completed: {result}")
        return result
        
    except Exception as e:
        logger.error(f"Enhanced SIGMA rule generation task failed: {e}")
        self.update_state(
            state='FAILURE',
            meta={
                'stage': 'error',
                'progress': 0,
                'message': f'Task failed: {str(e)}',
                'error': str(e)
            }
        )
        raise
    finally:
        db_session.close()