Add optional --timestamp and --category parameters

entrypoint.py

@@ -1,13 +1,20 @@
 #!/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'}
@@ -16,6 +23,15 @@ def close_issue(issue: dict, headers: Headers):
         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,
@@ -26,19 +42,26 @@ def create_new_issues(
         vulnerabilities: list,
         open_issues: OpenIssues,
         url: str,
-        headers: Headers,
+        headers: Headers
 ) -> set[str]:
-    found_issues = set()
+    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'])
-        found_issues.add(title)
+        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']],
+            '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()
@@ -49,12 +72,13 @@ def create_new_issues(
                 '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 found_issues
+    return issues_found_in_current_run
 
 
 def close_old_issues(
@@ -65,7 +89,13 @@ def close_old_issues(
     for title, open_issue_list in open_issues.items():
         if title not in found_issues:
             for open_issue in open_issue_list:
-                close_issue(open_issue, headers)
+                # 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:
@@ -78,6 +108,16 @@ def fetch_open_issues(url: str, headers: Headers) -> OpenIssues:
         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,
@@ -96,10 +136,17 @@ def fetch_open_issues(url: str, headers: Headers) -> OpenIssues:
             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:
@@ -113,8 +160,15 @@ def fetch_open_issues(url: str, headers: Headers) -> OpenIssues:
 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)
 
@@ -133,10 +187,10 @@ def main():
     open_issues = fetch_open_issues(url, headers)
     print('')
     print('=== Creating new issues ===')
-    found_issues = create_new_issues(vulnerabilities, open_issues, url, headers)
+    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, headers)
+    close_old_issues(open_issues, found_issues_in_current_run, headers)
 
 
 if __name__ == '__main__':