Merge ubuntu-cve-tracker:usn-publishing into ubuntu-cve-tracker:master

[Mark Morlino]

* improved error handling
* add ability to remove usns
* add option to upsert
* add option to control behavior during batch operations (exit on fail or keep going)

[Alex Murray (alexmurray)]

LGTM - I tested this during the publication of USN 4405-1 - https://ubuntu.com/security/notices/USN-4405-1 - and it seemed to work ok.

review: Approve

diff --git a/scripts/convert-pickle.py b/scripts/convert-pickle.py
index a3133b7..88258de 100755
--- a/scripts/convert-pickle.py
+++ b/scripts/convert-pickle.py
@@ -60,7 +60,7 @@ if __name__ == "__main__":
     parser.add_option("-i", "--input-file", dest="infile", help="pickle data file", metavar="FILE")
     parser.add_option("-o", "--output-file", dest="outfile", help="target json data file", metavar="FILE")
     parser.add_option("-p", "--prefix", dest="prefix",
-        help="prefix for value of id field. eg: 'USN-'", default=None, metavar="FILE")
+        help="prefix for value of id field. eg: 'USN-'", default=None)
     (opt, args) = parser.parse_args()
 
     if not opt.infile:
diff --git a/scripts/publish-usn-to-website b/scripts/publish-usn-to-website
index 33f63fd..20d2b58 100755
--- a/scripts/publish-usn-to-website
+++ b/scripts/publish-usn-to-website
@@ -77,19 +77,33 @@ push_to_website_api() {
 	TMP_PICKLE=$(mktemp -u)
 	TMP_JSON=$(mktemp -u)
 
+	# shellcheck disable=SC2064
+	trap "rm -f \"$TMP_YAML\" \"$TMP_PICKLE\" \"$TMP_JSON\"" EXIT
+
 	# export the single USN from the pickle to a yaml file
-	"$usn_tool"/usn.py --db "$database" --export "$usn" >"$TMP_YAML"
+	if ! "$usn_tool"/usn.py --db "$database" --export "$usn" >"$TMP_YAML" ; then 
+		perr "Failed to export $USN from pickle to yaml"
+		exit 1
+	fi
 
 	# import the single USN from the yaml into a new pickle file
-	"$usn_tool"/usn.py --db "$TMP_PICKLE" --import < "$TMP_YAML"
+	if ! "$usn_tool"/usn.py --db "$TMP_PICKLE" --import < "$TMP_YAML" ; then 
+		perr "Failed to create temporary pickle from $USN yaml"
+		exit 1
+	fi
 
 	# convert the single USN pickle to a single USN json file
-	"$UCT"/scripts/convert-pickle.py --input-file "$TMP_PICKLE" --output-file "$TMP_JSON" --prefix "USN-"
+	if ! "$UCT"/scripts/convert-pickle.py --input-file "$TMP_PICKLE" --output-file "$TMP_JSON" --prefix "USN-" ; then 
+		perr "Failed to convert $USN pickle to json"
+		exit 1
+	fi
 
 	# post the single USN json file to the website API
-	"$UCT"/scripts/publish-usn-to-website-api.py --action "$action" --json "$TMP_JSON"
+	if ! "$UCT"/scripts/publish-usn-to-website-api.py --action "$action" --json "$TMP_JSON" --stop --no-upsert; then 
+		perr "Failed to publish $USN pickle to website API"
+		exit 1
+	fi
 
-	rm -f "$TMP_YAML" "$TMP_PICKLE" "$TMP_JSON"
 }
 
 # $1: The action string (add, update, or remove)
diff --git a/scripts/publish-usn-to-website-api.py b/scripts/publish-usn-to-website-api.py
index cd3c92f..1185739 100755
--- a/scripts/publish-usn-to-website-api.py
+++ b/scripts/publish-usn-to-website-api.py
@@ -5,15 +5,23 @@
 #
 # initial upload of all USNs can be done as follows
 # cd $UCT && ./scripts/fetch-db database.pickle.bz2
-# ./scripts/convert-pickle.py --input database.pickle --output database.json
+# ./scripts/convert-pickle.py --input database.pickle --output database.json --prefix 'USN-'
 # ./scripts/publish-usn-to-website-api.py --json ./database.json --action add
+#
+# to test with a single USN you can use these commands:
+# export USN-9999
+# ../usn-tool/usn.py --db database.pickle --export $USN > $USN.yaml
+# ../usn-tool/usn.py --db $USN.pickle --import < $USN.yaml
+# ./scripts/convert-pickle.py --input-file $USN.pickle --output-file $USN.json --prefix USN- 
+# ./scripts/publish-usn-to-website-api.py --action add --json $USN.json 
+
 
 # Standard library
 import argparse
 import json
 import sys
 import os
-#import re
+import re
 from datetime import datetime
 from http.cookiejar import MozillaCookieJar
 
@@ -66,45 +74,60 @@ def guess_binary_links(binary, version, sources):
 
 parser = argparse.ArgumentParser(description="CLI to post USNs",)
 parser.add_argument(
-    "--json", action="store", type=str,
-    required=True,
-    help='path to json file',
-    )
-parser.add_argument(
     "--action", action="store", type=str,
     default="add",
     choices=['add', 'update', 'remove'],
-    help='api action to perform',
+    help='API action to perform',
+    )
+parser.add_argument(
+    "--json", action="store", type=str,
+    help='Path to json file. Required for --action add/update. Can be used instead of --usns for --action remove',
+    )
+parser.add_argument(
+    "--usns", action="store", nargs='+',
+    help='Can be used instead of --json with --action remove',
     )
 parser.add_argument(
     "--endpoint", action="store", type=str,
     default="https://ubuntu.com/security/notices",
-    help='api endpoint url',
+    help='API endpoint url. default is https://ubuntu.com/security/notices',
+    )
+parser.add_argument(
+    "--stop", action="store_true",
+    help="Exit after non-200 status. If upserting, exit after both attempted add/update operations fail",
+    )
+parser.add_argument(
+    "--no-upsert", action="store_true",
+    help="No update after non-200 status on add, No add after non-200 status on update",
     )
-#parser.add_argument(
-    #"--errors", action="store", type=str,
-    #default="exit",
-    #choices=['continue', 'exit'],
-    #help="what to do when the api returns a non-200 status",
-    #)
 
 args = parser.parse_args()
 
-if args.action == 'add':
-    http_method = 'POST'
-elif args.action == 'update':
-    http_method = 'PUT'
-elif args.action == 'remove':
-    sys.exit(f"Error: {args.action} action not implemented yet")
-else:
-    sys.exit(f"Error: {args.action} action not implemented yet")
+if args.json and not os.path.exists(args.json):
+    sys.exit(f"Error: {args.json} not found")
+if (args.action == 'add' or args.action == 'update') and not args.json:
+    sys.exit(f"Error: --action {args.action} requires --json")
+if args.action == 'remove' and not args.json and not args.usns:
+    sys.exit("Error: --action remove requires --json or --usns")
 
 client = httpbakery.Client(cookies=MozillaCookieJar(".login"))
 
 # Make a first call to make sure we are logged in
-response = client.request(http_method, url=args.endpoint)
-
-
+response = client.request('GET', url=args.endpoint)
+
+# if we have a list of USNs instead of a json file
+# we can just remove/DELETE them
+if args.action == 'remove' and args.usns:
+    http_method = 'DELETE'
+    for usn in args.usns:
+        endpoint = f"{args.endpoint}/{usn}"
+        response = client.request( http_method, url=endpoint, json=None)
+        print(response, response.text[0:60])
+        if args.stop and not re.match(r'^<Response \[2..\]>$',str(response)):
+            sys.exit(1)
+    sys.exit(0)
+
+# read in the json
 with open(args.json) as usn_json:
     payload = json.load(usn_json).items()
 
@@ -146,30 +169,50 @@ with open(args.json) as usn_json:
             else:
                 references.append(reference)
 
+        # Build json payload
+        json_data = {
+            "id": notice["id"],
+            "description": notice["description"],
+            "references": references,
+            "cves": cves,
+            "release_packages": release_packages,
+            "title": notice["title"],
+            "published": datetime.fromtimestamp(
+                notice["timestamp"]
+            ).isoformat(),
+            "summary": notice["summary"],
+            "instructions": notice.get("action"),
+        }
+        
         # Build endpoint
-        endpoint = args.endpoint
-        if http_method == "PUT":
+        if args.action == 'add':
+            upsert = not args.no_upsert
+            http_method = 'POST'
+            endpoint = args.endpoint
+        elif args.action == 'update':
+            upsert = not args.no_upsert
+            http_method = 'PUT'
             endpoint = f"{args.endpoint}/{notice['id']}"
-
-        response = client.request(
-            http_method,
-            url=endpoint,
-            json={
-                "id": notice["id"],
-                "description": notice["description"],
-                "references": references,
-                "cves": cves,
-                "release_packages": release_packages,
-                "title": notice["title"],
-                "published": datetime.fromtimestamp(
-                    notice["timestamp"]
-                ).isoformat(),
-                "summary": notice["summary"],
-                "instructions": notice.get("action"),
-            },
-        )
-
-        print(response, response.text)
-
-        #if args.errors == 'exit' and not re.match("2\d\d",str(response)):
-            #sys.exit(1)
+        elif args.action == 'remove':
+            upsert = False
+            http_method = 'DELETE'
+            endpoint = f"{args.endpoint}/{notice['id']}"
+            json_data = None
+
+        response = client.request( http_method, url=endpoint, json=json_data)
+        print(response, response.text[0:60])
+
+        if upsert and not re.match(r'^<Response \[2..\]>$',str(response)):
+            if args.action == 'add':
+                print("add failed, trying update")
+                http_method = 'PUT'
+                endpoint = f"{args.endpoint}/{notice['id']}"
+            elif args.action == 'update':
+                print("update failed, trying add")
+                http_method = 'POST'
+                endpoint = args.endpoint
+            response = client.request( http_method, url=endpoint, json=json_data)
+            print(response, response.text[0:60])
+
+        if args.stop and not re.match(r'^<Response \[2..\]>$',str(response)):
+            sys.exit(1)