fix w shedulerze
This commit is contained in:
Mateusz Gruszczyński
150
scheduler.py
150
scheduler.py
@@ -2,8 +2,8 @@
|
||||
|
||||
import time
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from cve_handler import update_all_vendors
|
||||
from datetime import datetime, timedelta
|
||||
from cve_handler import CVEHandler
|
||||
import config
|
||||
|
||||
logging.basicConfig(
|
||||
@@ -12,29 +12,153 @@ logging.basicConfig(
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
def run_scheduler():
|
||||
logger.info("CVE Update Scheduler started")
|
||||
logger.info(f"Update interval: {config.UPDATE_INTERVAL_HOURS} hours")
|
||||
logger.info(f"Auto-update enabled: {config.ENABLE_AUTO_UPDATE}")
|
||||
|
||||
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:
|
||||
logger.info(f"Starting scheduled update at {datetime.now()}")
|
||||
updated, failed = update_all_vendors()
|
||||
logger.info(f"Update completed: {updated} successful, {failed} failed")
|
||||
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 = config.UPDATE_INTERVAL_HOURS * 3600
|
||||
next_update = datetime.now().timestamp() + sleep_seconds
|
||||
logger.info(f"Next update scheduled at {datetime.fromtimestamp(next_update)}")
|
||||
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__':
|
||||
run_scheduler()
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user