import json
import logging
import time

from collections import defaultdict

import features

from app import secscan_notification_queue, secscan_api
from data import model
from data.database import (Image, ImageStorage, ExternalNotificationEvent,
                           Repository, RepositoryNotification, RepositoryTag)
from endpoints.notificationhelper import spawn_notification
from workers.queueworker import QueueWorker


logger = logging.getLogger(__name__)


class SecurityNotificationWorker(QueueWorker):
  def process_queue_item(self, queueitem):
    data = json.loads(queueitem.body)

    cve_id = data['Name']
    vulnerability = data['Content']['Vulnerability']
    priority = vulnerability['Priority']

    # Lookup the external event for when we have vulnerabilities.
    event = ExternalNotificationEvent.get(name='vulnerability_found')

    # For each layer, retrieving the matching tags and join with repository to determine which
    # require new notifications.
    tag_map = defaultdict(set)
    repository_map = {}

    # Find all tags that contain the layer(s) introducing the vulnerability.
    content = data['Content']
    layer_ids = content.get('NewIntroducingLayersIDs', content.get('IntroducingLayersIDs', []))
    for layer_id in layer_ids:
      (docker_image_id, storage_uuid) = layer_id.split('.', 2)
      tags = model.tag.get_matching_tags(docker_image_id, storage_uuid, RepositoryTag,
                                         Repository, Image, ImageStorage)

      # Additionally filter to tags only in repositories that have the event setup.
      matching = list(tags
                      .switch(RepositoryTag)
                      .join(Repository)
                      .join(RepositoryNotification)
                      .where(RepositoryNotification.event == event))

      check_map = {}
      for tag in matching:
        # Verify that the tag's root image has the vulnerability.
        tag_layer_id = '%s.%s' % (tag.image.docker_image_id, tag.image.storage.uuid)
        logger.debug('Checking if layer %s is vulnerable to %s', tag_layer_id, cve_id)

        if not tag_layer_id in check_map:
          is_vulerable = secscan_api.check_layer_vulnerable(tag_layer_id, cve_id)
          check_map[tag_layer_id] = is_vulerable

        logger.debug('Result of layer %s is vulnerable to %s check: %s', tag_layer_id, cve_id,
                     check_map[tag_layer_id])

        if check_map[tag_layer_id]:
          # Add the vulnerable tag to the list.
          tag_map[tag.repository_id].add(tag.name)
          repository_map[tag.repository_id] = tag.repository

    # For each of the tags found, issue a notification.
    for repository_id in tag_map:
      tags = tag_map[repository_id]
      event_data = {
        'tags': list(tags),
        'vulnerability': {
          'id': data['Name'],
          'description': vulnerability['Description'],
          'link': vulnerability['Link'],
          'priority': priority,
        },
      }

      spawn_notification(repository_map[repository_id], 'vulnerability_found', event_data)


if __name__ == '__main__':
  if not features.SECURITY_SCANNER:
    logger.debug('Security scanner disabled; skipping SecurityNotificationWorker')
    while True:
      time.sleep(100000)

  worker = SecurityNotificationWorker(secscan_notification_queue, poll_period_seconds=30,
                                      reservation_seconds=30, retry_after_seconds=30)
  worker.start()