entrypoint.py

#!/usr/bin/env python3

import argparse
import json
from datetime import datetime
import random
import time

import requests
import os

Headers = dict[str, str]
OpenIssues = dict[str, list[dict]]
timestamp=None
category=None
pause_lowerbound_sec=10
pause_upperbound_sec=20

def close_issue(issue: dict, headers: Headers):
    payload = {'state': 'closed'}
    r = requests.patch(issue['url'], data=json.dumps(payload), headers=headers)
    if r.ok:
        print('Closed issue {}#{}'.format(
            issue['title'], issue['number'],
        ))
    elif r.headers.get('retry-after'):
        retry_after = int(r.headers['retry-after'])
        print('Rate limited, waiting for {} seconds'.format(retry_after))
        time.sleep(retry_after)
        close_issue(issue, headers)
    elif r.headers.get('x-ratelimit-remaining') == '0':
        print('Rate limited, waiting for 60 seconds')
        time.sleep(r.headers.get('x-ratelimit-reset', 60))
        close_issue(issue, headers)
    else:
        print('ERROR: cannot close issue {}#{} (status code {}): {}'.format(
            issue['title'], issue['number'], r.status_code, r.text,
        ))


def create_new_issues(
        vulnerabilities: list,
        open_issues: OpenIssues,
        url: str,
        headers: Headers
) -> set[str]:
    issues_found_in_current_run = set()

    for vulnerability in vulnerabilities:
        time.sleep(random.randint(pause_lowerbound_sec, pause_upperbound_sec)) # to prevent rate limiting
        title = '{} {}'.format(vulnerability['title'], vulnerability['id'])
        issues_found_in_current_run.add(title)
        issues_found_in_current_run.add(title)
        if title in open_issues:
            continue
        payload = {
            'title': title,
            'body': vulnerability['description'],
            'labels': [
                'Snyk', vulnerability['severity'], f"Snyk-timestamp:{timestamp}",
            ],
        }
        if category is not None:
            payload['labels'].append(category)
        r = requests.post(url, data=json.dumps(payload), headers=headers)
        if r.ok:
            new_issue = r.json()
            print('Created issue {}#{}: status code {}'.format(
                title, new_issue['number'], r.status_code,
            ))
            open_issues[title] = [{
                'title': title,
                'url': new_issue['url'],
                'number': new_issue['number'],
                'timestamp': timestamp,
            }]
        else:
            print('Failed to create issue {} (status code {}): {}'.format(
                title, r.status_code, r.text,
            ))
    return issues_found_in_current_run


def close_old_issues(
        open_issues: OpenIssues,
        found_issues: set[str],
        headers: Headers,
):
    for title, open_issue_list in open_issues.items():
        if title not in found_issues:
            for open_issue in open_issue_list:
                # Only close issues when the open issue is not of the same meta Snyk run
                # i.e. does not have the same Snyk-timestamp label (only if provided).
                time.sleep(random.randint(pause_lowerbound_sec, pause_upperbound_sec)) # to prevent rate limiting
                if (timestamp is not None and
                        timestamp is not open_issue.get('timestamp')):
                    print(f"Closing issue {open_issue['title']} - #{open_issue['number']}")
                    close_issue(open_issue, headers)


def fetch_open_issues(url: str, headers: Headers) -> OpenIssues:
    open_issues = dict()

    next_url = url + '?labels=Snyk&per_page=100'
    num_tries = 0
    while next_url:
        print('Fetching {}'.format(next_url))
        r = requests.get(next_url, headers=headers)
        if r.ok:
            num_tries = 0
        # Handle rate limiting (see https://developer.github.com/v3/#rate-limiting)
        elif r.headers.get('retry-after'):
            retry_after = int(r.headers['retry-after'])
            print('Rate limited, waiting for {} seconds'.format(retry_after))
            time.sleep(retry_after)
            continue
        elif r.headers.get('x-ratelimit-remaining') == '0':
            print('Rate limited, waiting for 60 seconds')
            time.sleep(r.headers.get('x-ratelimit-reset', 60))
            continue
        elif num_tries < 3:
            print('Failed to fetch issues {}, retrying (status code {}): {}'.format(
                next_url, r.status_code, r.text,
            ))
            num_tries += 1
            continue
        else:
            raise IOError('Failed to fetch issues {} (status code {}): {}'.format(
                next_url, r.status_code, r.text,
            ))

        for open_issue in r.json():
            # Skip PR "issues"
            if 'pull_request' in open_issue:
                continue
            title = open_issue['title']
            if title not in open_issues:
                open_issues[title] = []
            # Get the timestamp label if it exists
            issue_timestamp = None
            for label in open_issue['labels']:
                if label['name'].startswith('Snyk-timestamp:'):
                    issue_timestamp = label['name'].split(':')[1]
                    break
            open_issues[title].append({
                'title': open_issue['title'],
                'url': open_issue['url'],
                'number': open_issue['number'],
                'timestamp': issue_timestamp,
            })
        next_link = r.links.get('next')
        if next_link:
            next_url = next_link['url']
        else:
            next_url = None

    return open_issues


def main():
    parser = argparse.ArgumentParser(description="Parse the report file")
    parser.add_argument("report_file", type=str, help="a name of the report file")
    parser.add_argument("-t", "--timestamp", type=str, help="a timestamp of the run. When passed, will cause combine results of multiple runs when the runs have the same timestamp.", default=datetime.now(), required=False)
    parser.add_argument("-c", "--category", type=str, help="a category label to add to the issue.", required=False)
    args = parser.parse_args()

    global timestamp
    timestamp = args.timestamp
    global category
    category = args.category

    with open(args.report_file) as f:
        report = json.load(f)

    if isinstance(report, list):
        vulnerabilities = [v for obj in report for v in obj['vulnerabilities']]
    else:
        vulnerabilities = [v for v in report['vulnerabilities']]

    github_repository = os.environ['GITHUB_REPOSITORY']
    url = 'https://api.github.com/repos/{}/issues'.format(github_repository)
    token = os.environ['TOKEN']
    headers = {'Authorization': 'token {}'.format(token),
               'Accept': 'application/vnd.github.v3+json'}

    print('=== Fetching open issues ===')
    open_issues = fetch_open_issues(url, headers)
    print('')
    print('=== Creating new issues ===')
    found_issues_in_current_run = create_new_issues(vulnerabilities, open_issues, url, headers)
    print('')
    print('=== Closing old issues ===')
    close_old_issues(open_issues, found_issues_in_current_run, headers)


if __name__ == '__main__':
    main()