process-cve-records: add --work-dir argument

And switch to argparse for argument processing.

1 files changed, 17 insertions, 15 deletions

diff --git a/bin/process-cve-records b/bin/process-cve-records
index 135e19c27e..cf6ce184ad 100755
--- a/bin/process-cve-records
+++ b/bin/process-cve-records
@@ -7,10 +7,10 @@
 #
 # Copyright © 2023 Emilio Pozuelo Monfort <pochu@debian.org>
 
+import argparse
 import io
 import json
 import os
-import sys
 import zipfile
 
 import requests
@@ -20,10 +20,8 @@ from sectracker import parsers
 
 CVE_ZIPFILE = 'https://github.com/CVEProject/cvelistV5/archive/refs/heads/main.zip'
 
-debug_enabled = False
-
 def debug(m):
-    if debug_enabled:
+    if args.verbose:
         print(m)
 
 
@@ -126,8 +124,15 @@ def download_zip_file():
     b = io.BytesIO(r.content)
     process_zip_file(b)
 
+default_workdir = os.path.join(os.path.dirname(os.path.dirname(__file__)))
+
+parser = argparse.ArgumentParser(description='Update CVE list with MITRE CVE records')
+parser.add_argument('-v', '--verbose', action="store_true", help='enable verbose messages')
+parser.add_argument('--work-dir', help='path to security-tracker repo (default: relative to the script)', default=default_workdir)
+parser.add_argument('file', nargs='?', help='file to process, or download records from MITRE if not specified')
+args = parser.parse_args()
 
-main_list = os.path.dirname(__file__) + '/../data/CVE/list'
+main_list = args.work_dir + '/data/CVE/list'
 
 debug("reading cve file")
 cves = parsers.cvelist(main_list)
@@ -135,20 +140,17 @@ debug("finished reading cve file")
 
 cve_dir = { cve.header.name: cve for cve in cves }
 
-if len(sys.argv) == 1:
+if not args.file:
     # no argument, we download the CVE db
     download_zip_file()
-elif sys.argv[1].endswith('.json'):
-    record_file = sys.argv[1]
-    debug("processing record " + record_file)
-    process_record_filename(record_file)
+elif args.file.endswith('.json'):
+    debug("processing record " + args.file)
+    process_record_filename(args.file)
     debug("record processed")
-elif sys.argv[1].endswith('.zip'):
-    zip_file = sys.argv[1]
-    process_zip_file(zip_file)
+elif args.file.endswith('.zip'):
+    process_zip_file(args.file)
 else:
-    record_dir = sys.argv[1]
-    process_record_dir(record_dir)
+    process_record_dir(args.file)
 
 # write CVE file back
 with open(main_list, 'w') as f: