cve_monitor/scheduler.py

#!/usr/bin/env python3

import time
import logging
from datetime import datetime, timedelta
from cve_handler import CVEHandler
import config

logging.basicConfig(
    level=config.LOG_LEVEL,
    format=config.LOG_FORMAT
)
logger = logging.getLogger(__name__)


def incremental_update():

    logger.info("=" * 80)
    logger.info(f"Starting incremental update at {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    logger.info("=" * 80)
    
    try:
        handler = CVEHandler()
        days_back = getattr(config, 'SCHEDULER_UPDATE_DAYS', 7)
        
        updated_count = 0
        failed_count = 0
        total_new_cves = 0
        
        for vendor in config.VENDORS:
            if not vendor.get('use_cpe'):
                logger.debug(f"Skipping {vendor['name']} (CPE disabled)")
                continue
            
            try:
                logger.info(f"Updating {vendor['name']}...")
                
                new_cves = handler.fetch_cve_from_nvd(
                    keywords=vendor.get('keywords', []),
                    vendor_config=vendor,
                    days_back=days_back
                )
                
                if new_cves:
                    handler.save_cves_to_cache(new_cves, vendor['code'])
                    logger.info(f"  ✓ {vendor['name']}: Added {len(new_cves)} new CVEs")
                    total_new_cves += len(new_cves)
                    updated_count += 1
                else:
                    logger.info(f"  ✓ {vendor['name']}: No new CVEs")
                    updated_count += 1
                
            except Exception as e:
                logger.error(f"  ✗ {vendor['name']}: {e}")
                failed_count += 1
        
        logger.info("=" * 80)
        logger.info(f"✓ Incremental update completed")
        logger.info(f"  Vendors updated: {updated_count}")
        logger.info(f"  Vendors failed: {failed_count}")
        logger.info(f"  New CVEs added: {total_new_cves}")
        logger.info("=" * 80)
        
        return updated_count, failed_count
        
    except Exception as e:
        logger.error(f"✗ Incremental update failed: {e}", exc_info=True)
        return 0, len(config.VENDORS)


def initial_full_scan():

    logger.info("=" * 80)
    logger.info("INITIAL FULL SCAN")
    logger.info("=" * 80)
    logger.info(f"This will scan {config.INITIAL_LOOKBACK_DAYS} days back for all vendors")
    
    try:
        from cve_handler import update_all_vendors
        updated, failed = update_all_vendors(force=True)
        
        logger.info("=" * 80)
        logger.info(f"✓ Initial scan completed: {updated} updated, {failed} failed")
        logger.info("=" * 80)
        
        return updated, failed
        
    except Exception as e:
        logger.error(f"✗ Initial scan failed: {e}", exc_info=True)
        return 0, len(config.VENDORS)


def run_scheduler():
    logger.info("=" * 80)
    logger.info("CVE Monitor Scheduler Starting")
    logger.info("=" * 80)
    
    update_interval_hours = config.UPDATE_INTERVAL_HOURS
    scheduler_update_days = getattr(config, 'SCHEDULER_UPDATE_DAYS', 7)
    
    logger.info(f"Configuration:")
    logger.info(f"  Update interval: {update_interval_hours} hours")
    logger.info(f"  Incremental lookback: {scheduler_update_days} days")
    logger.info(f"  Initial lookback: {config.INITIAL_LOOKBACK_DAYS} days")
    logger.info(f"  Auto-update enabled: {config.ENABLE_AUTO_UPDATE}")
    
    if config.ENABLE_AUTO_UPDATE:
        logger.info("\nChecking database status...")
        
        try:
            handler = CVEHandler()
            
            with handler.get_db_connection() as conn:
                cursor = conn.cursor()
                cursor.execute("SELECT COUNT(*) FROM cve_cache")
                total_cves = cursor.fetchone()[0]
            
            if total_cves == 0:
                logger.info(f"  Database is EMPTY")
                logger.info(f"Starting initial full scan ({config.INITIAL_LOOKBACK_DAYS} days)...")
                logger.info(f"This may take 30-60 minutes depending on API rate limits\n")
                
                initial_full_scan()
                
            else:
                logger.info(f"✓ Database has {total_cves} CVEs")
                logger.info(f"Skipping initial scan, starting incremental updates\n")
                logger.info("Performing startup incremental update...")
                incremental_update()
                
        except Exception as e:
            logger.error(f"Error checking database: {e}", exc_info=True)
    
    logger.info(f"\nWaiting 2 minutes before starting scheduled updates...")
    time.sleep(120)
    
    while True:
        if config.ENABLE_AUTO_UPDATE:
            try:
                updated, failed = incremental_update()
                
            except Exception as e:
                logger.error(f"Error in scheduled update: {e}", exc_info=True)
        else:
            logger.debug("Auto-update is disabled, skipping")
        
        sleep_seconds = update_interval_hours * 3600
        next_update = datetime.now() + timedelta(seconds=sleep_seconds)
        
        logger.info(f"\n{'=' * 80}")
        logger.info(f"Next incremental update at: {next_update.strftime('%Y-%m-%d %H:%M:%S')}")
        logger.info(f"Sleeping for {update_interval_hours} hours...")
        logger.info(f"{'=' * 80}\n")
        
        time.sleep(sleep_seconds)


if __name__ == '__main__':
    try:
        run_scheduler()
    except KeyboardInterrupt:
        logger.info("\n\nScheduler stopped by user")
    except Exception as e:
        logger.error(f"Scheduler crashed: {e}", exc_info=True)