snyk_feed.py

#!/usr/bin/env python3
# Copyright © 2020 Red Hat Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Author: Yusuf Zainee <yzainee@redhat.com>
#

"""Script which fetches the feed from snyk and stores it in S3 bucket."""
import gzip
import os
import urllib
import requests
import logging
import json
from datetime import datetime
from graph_snyk_cve_sync import SnykCveSync
from helper import Helper
from status_check import StatusCheck

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# api fields response key mapping
# Mapping from old key to new key
old_new_keys_mapping = {
    'creationTime': 'published',
    'credit': 'credits',
    'cves': 'cve_ids',
    'cvssScore': 'cvss_v3_base_score',
    'cvssV3': 'cvss_v3_vector',
    'cwes': 'cwe_ids',
    'descriptionOverview': 'descirption_overview',
    'descriptionRemediation': 'description_remediation',
    'disclosureTime': 'disclosed',
    'exploit': 'exploit_code_maturity',
    'fixable': 'is_fixable',
    'functions': 'vulnerable_functions',
    'functions_new': 'vulnerable_functions',
    'id': 'snyk_id',
    'initiallyFixedIn': 'initially_fixed_in_versions',
    'language': 'ecosystem',
    'malicious': 'is_malicious',
    'modificationTime': 'modified',
    'patchExists': 'is_fixable',
    'publicationTime': 'published',
    'registry': 'package_repository_url',
    'socialTrendAlert': 'is_social_media_trending',
    'url': 'snyk_advisory_url',
    'vulnerableVersions': 'vulnerable_versions',
    'vulnerableHashRanges': 'vulnerable_hash_ranges',
    'vulnerableHashes': 'vulnerable_hashes'
}


# in order to achieve old format json
def _arrange_ecosystems(snyk_feed_json):
    java_json = []
    javascript_json = []
    python_json = []
    golang_json = []
    # to reduce size of data not considering
    # data other than the ecosystem we support
    for row in snyk_feed_json:
        for item in row:
            if item == "language":
                if row[item] == "java":
                    java_json.append(row)
                elif row[item] == "js":
                    javascript_json.append(row)
                elif row[item] == "python":
                    python_json.append(row)
                if row[item] == "golang":
                    golang_json.append(row)
    logger.info("Combining all the ecosystems together and ignoring the ecosystems we don't support")
    formatted_json = {"java": java_json, "js": javascript_json, "python": python_json,
                      "golang": golang_json}

    return formatted_json


def _convert_ndjson_to_json(ndjson_content):
    json_content = []
    for ndjson_line in ndjson_content.splitlines():
        if not ndjson_line.strip():
            continue  # ignore empty lines
        json_line = json.loads(ndjson_line)
        json_content.append(json_line)
    return json_content


def _rename_new_keys_to_old_keys(json_content):
    for row in json_content:
        for k, v in old_new_keys_mapping.items():
            for new_name in row:
                if v == new_name:
                    row[k] = row.pop(new_name)
    return json_content


def _download_file(url):
    # Download archive
    try:
        # Read the file inside the .gz archive located at url
        with urllib.request.urlopen(url) as response:
            with gzip.GzipFile(fileobj=response) as uncompressed:
                ndjson_content = uncompressed.read()
                logger.info("Extracting the downloaded zip file")
                logger.info("Reading file content")
                del response, uncompressed
                return ndjson_content
    except Exception as e:
        print(e)
        return 1


class SnykDataFetcher:
    """Snyk data class to fetch the information from Snyk."""

    def __init__(self):
        """Init function to add default values."""
        self.helper = Helper()
        self.snyk_data = {}
        self.ORG_ID = os.environ.get('SNYK_ORG_ID', '')
        self.URL = os.environ.get('SNYK_URL',
                                  'https://api.snyk.io/rest/orgs/{}/intel_feed/application_premium?version=2022-09-29').format(
            self.ORG_ID)
        self.TOKEN = os.environ.get('SNYK_TOKEN', '')
        self.DISABLE_SNYK_SYNC_OPERATION = os.environ.get('DISABLE_SNYK_SYNC_OPERATION', 'false') \
                                               .lower() in ('1', 'yes', 'true')

    def _validate(self):
        """Validate method to see all the params are set.

        :return True if all params are valid. Raise exception otherwise & halt proceedings.
        """
        StatusCheck().status_check()
        try:
            assert self.ORG_ID
            assert self.URL
            assert self.TOKEN
        except AssertionError:
            raise ValueError('Snyk Feed fetch aborted. Either ORG_ID, URL or TOKEN missing')
        logger.info("Snyk credentials validated.")
        return True

    def _fetch(self):
        """Fetch function to fetch the feed from Snyk.

        :return True if the operation to fetch & store is successful. Exception otherwise.
        """
        headers = {
            'Authorization': ' '.join(['token', self.TOKEN])
        }
        try:
            res = requests.get(self.URL, headers=headers)

            logger.info("status_code : {}".format(res.status_code))
        except Exception as e:
            logger.error('Unable to get snyk feed. Reason: %r' % e)
            return False
        response = json.loads(res.text)
        logger.info("Snyk feed successfully fetched.")

        data_url = response['data']['url']
        logger.info("Downloading snyk feed zip file...")
        ndjson_snyk_feed_content = _download_file(data_url)

        logger.info("Converting the ndjson snyk feed data to json format...")
        json_snyk_feed_content = _convert_ndjson_to_json(ndjson_snyk_feed_content)
        # since it is no longer required
        del ndjson_snyk_feed_content

        logger.info("Renaming new fields to old fields")
        formatted_snyk_feed_json = _rename_new_keys_to_old_keys(json_snyk_feed_content)
        # since it is no longer required
        del json_snyk_feed_content

        logger.info("Arranging ecosystems in an order")
        formatted_snyk_feed_json = _arrange_ecosystems(formatted_snyk_feed_json)

        self.snyk_data = formatted_snyk_feed_json
        logger.info("Data successfully loaded")
        return True

    def _run_snyk_ingestion(self):
        """Run the process of parsing and ingesting snyk data."""
        logger.info("Triggering Snyk ingestion.")
        return SnykCveSync(self.snyk_data).run_snyk_sync()

    def _store_data(self):
        """Store the snyk data into the respective S3 bucket."""
        self.helper.store_json_content(self.snyk_data, "snyk-feed/latest-feed.json")
        logger.info("Snyk feed saved in S3.")

    def _disable_snyk_run(self, date):
        """Enable or disable snyk run.

        :return True if this needs to be run forcefully or if its saturday midnight.
        False otherwise.
        """
        run_time = self.helper.ingestion_run_time()
        logger.info("Disable Snyk Sync {}".format(self.DISABLE_SNYK_SYNC_OPERATION))
        logger.info("Force Snyk Run {}".format(self.helper.force_run_ingestion()))
        logger.info("Current Hour {}".format(date.strftime('%H')))
        logger.info("Full/Delta Ingestion Run Time {}".format(run_time))
        if self.DISABLE_SNYK_SYNC_OPERATION:
            return True
        force = self.helper.force_run_ingestion()
        if force:
            return False
        if run_time != "none":
            return date.strftime('%H') != self.helper.ingestion_run_time()
        return False

    def run_snyk_fetch(self):
        """Entry function for the snyk fetch feed.

        :return Success message for completion of task.
        """
        logger.info("Snyk Feed Fetch Function".center(50, '-'))
        self._validate()
        utc_now = datetime.utcnow()
        if self._disable_snyk_run(utc_now):
            logger.info("Snyk feed wont be fetched as the operations are disabled at the moment.")
            return
        logger.info("Snyk Feed Fetch Begins")
        if self._fetch():
            self._store_data()
            self._run_snyk_ingestion()
        return "Success"