Merge ~mdeslaur/ubuntu-cve-tracker:convert-pickle-py3 into ubuntu-cve-tracker:master

[Marc Deslauriers]

commit 881703b3dc9ba3eabcab51fc6105c2d38933e55a
Author: Marc Deslauriers <email address hidden>
Date: Tue Mar 12 09:57:49 2024 -0400

    scripts/convert-pickle.py: convert to python 3

commit 41d0f307086483417748c0327a21a0dd80a9ef3e
Author: Marc Deslauriers <email address hidden>
Date: Tue Mar 12 09:52:48 2024 -0400

    scripts/test_convert_pickle.py: add test for convert-pickle.py

This merge proposal converts convert-pickle.py to python 3 and gets rid of special handling for old python 2 versions that slows down execution.

8a188cd... by Marc Deslauriers on 2024-03-12

test_convert_pickle.py: clarify comment

[Alex Murray (alexmurray)]

LGTM - but one thing, can you please update .lauchpad.yaml to include your new test script? Thanks.

review: Approve

[Steve Beattie (sbeattie)]

On Tue, Mar 12, 2024 at 11:42:21PM -0000, Alex Murray wrote:
> Review: Approve
>
> LGTM - but one thing, can you please update .lauchpad.yaml to include your new test script? Thanks.

I'm planning to land my changes that adjust lpci / make check-python to
run pytest-3 against scripts/test_*.y so they will automatically get
picked up.

--
Steve Beattie
<email address hidden>

[Marc Deslauriers (mdeslaur)]

I've pushed this as-is as to not conflict with Steve's pending merge that will pick it up automatically.

diff --git a/scripts/convert-pickle.py b/scripts/convert-pickle.py
index 88258de..2eec39e 100755
--- a/scripts/convert-pickle.py
+++ b/scripts/convert-pickle.py
@@ -1,14 +1,13 @@
-#!/usr/bin/env python2
+#!/usr/bin/env python3
 
 # Author: Jamie Strandboge <[e-mail]>
-# Copyright (C) 2012 Canonical Ltd.
+# Copyright (C) 2012-2024 Canonical Ltd.
 #
 # This script is distributed under the terms and conditions of the GNU General
 # Public License, Version 2 or later. See http://www.gnu.org/copyleft/gpl.html
 # for details.
 
-import codecs
-import cPickle
+import pickle
 import json
 import optparse
 import os
@@ -18,7 +17,8 @@ import sys
 def save_database_json(database, db_filename):
     '''Save usn database'''
     filename = os.path.expanduser(db_filename)
-    json.dump(database, open(filename, 'w'), -1, encoding="utf-8")
+    with open(filename, 'w') as f:
+        json.dump(database, f)
 
 def convert_pickle_to_json(indb, outdb, prefix=None):
     '''Convert a pickle database into a json'''
@@ -26,31 +26,18 @@ def convert_pickle_to_json(indb, outdb, prefix=None):
     if not os.path.isfile(filename):
         return False
 
-    print >>sys.stderr, "INFO: Loading %s..." % (indb)
-    db = cPickle.load(open(indb))
+    print(f"INFO: Loading {indb}...", file=sys.stderr)
+    with open(indb, 'rb') as f:
+        db = pickle.load(f, encoding="utf-8")
 
-    # Older python pickle's have a bug that stores utf-8 data incorrectly.
-    # Account for that in our top level keys (db[k]['description'] is known to
-    # have this)
-    new_db = dict()
-    count = 0
-    for k in db.keys():
-        new_db[k] = dict()
-        for j in db[k].keys():
-            if prefix and j == 'id':
-                db[k][j] = "{}{}".format(prefix, db[k][j])
-            try:
-                json.dumps(db[k][j]) # if this fails, so will json.dump later
-                new_db[k][j] = db[k][j]
-            except:
-                count += 1
-                new_db[k][j] = db[k][j].decode("utf-8", "replace")
+    if prefix:
+        for k in db.keys():
+            for j in db[k].keys():
+                if j == 'id':
+                    db[k][j] = f"{prefix}{db[k][j]}"
 
-    if count > 0:
-        print >>sys.stderr, "WARN: performed %d pickle decode conversions" % count
-
-    print >>sys.stderr, "INFO: Saving %s..." % (outdb)
-    save_database_json(new_db, outdb)
+    print(f"INFO: Saving {outdb}...", file=sys.stderr)
+    save_database_json(db, outdb)
 
 #
 # main
@@ -64,16 +51,16 @@ if __name__ == "__main__":
     (opt, args) = parser.parse_args()
 
     if not opt.infile:
-        print >>sys.stderr, "Must specify --input-file"
+        print("Must specify --input-file", file=sys.stderr)
         sys.exit(1)
     elif not opt.outfile:
-        print >>sys.stderr, "Must specify --output-file"
+        print("Must specify --output-file", file=sys.stderr)
         sys.exit(1)
     elif not os.path.isfile(opt.infile):
-        print >>sys.stderr, "'%s' does not exist" % (opt.infile)
+        print(f"'{opt.infile}' does not exist", file=sys.stderr)
         sys.exit(1)
     elif os.path.exists(opt.outfile):
-        print >>sys.stderr, "'%s' already exists" % (opt.outfile)
+        print(f"'{opt.outfile}' already exists", file=sys.stderr)
         sys.exit(1)
 
     convert_pickle_to_json(opt.infile, opt.outfile, opt.prefix)
diff --git a/scripts/test_convert_pickle.py b/scripts/test_convert_pickle.py
 new file mode 100755
index 0000000..0fae6d9
--- /dev/null
+++ b/scripts/test_convert_pickle.py
@@ -0,0 +1,120 @@
+#!/usr/bin/env pytest-3
+# -*- coding: utf-8 -*-
+#
+# Author: Marc Deslauriers <[e-mail]>
+# Copyright (C) 2024 Canonical Ltd.
+#
+# This script is distributed under the terms and conditions of the GNU General
+# Public License, Version 3 or later. See http://www.gnu.org/copyleft/gpl.html
+# for details.
+#
+# Simple tests for convert-pickle.py
+#
+# The test pickle files contain USN-6686-1. This USN was selected randomly
+# because it was big and contained unicode characters.
+
+import pytest
+import subprocess
+import pickle
+import json
+import tempfile
+import os
+import unittest
+import shutil
+
+
+class TestConvertPickle(unittest.TestCase):
+    '''This tests the convert-pickle.py utility.'''
+
+    def setUp(self):
+        """Set up prior to each test_* function"""
+        self.tmpdir = tempfile.mkdtemp(prefix='test_convert_pickle-')
+
+    def tearDown(self):
+        """Clean up after each test_* function"""
+        if os.path.exists(self.tmpdir):
+            shutil.rmtree(self.tmpdir)
+
+    def _load_pickle(self, filename):
+        '''loads a pickle file'''
+        with open(filename, 'rb') as f:
+            db = pickle.load(f, encoding='utf-8')
+        return db
+
+    def _load_json(self, filename):
+        '''loads a json file'''
+        with open(filename, 'rb') as f:
+            db = json.load(f)
+        return db
+
+    def _load_text(self, filename):
+        '''loads a file as text'''
+        with open(filename, 'rb') as f:
+            text = f.read()
+        return text
+
+    def test_convert_pickle_v2(self):
+        '''test converting pickle v2 file to json'''
+        pickle_file = 'scripts/testfiles/database-protv2.pickle'
+        json_file = os.path.join(self.tmpdir, 'test.json')
+
+        rc = subprocess.call(["scripts/convert-pickle.py",
+                              "-i", pickle_file,
+                              "-o", json_file])
+        self.assertEqual(rc, 0)
+
+        pickle_db = self._load_pickle(pickle_file)
+        json_db = self._load_json(json_file)
+        json_text = self._load_text(json_file)
+
+        # Make sure the databases are the same
+        self.assertEqual(pickle_db, json_db)
+
+        # Make sure unicode is escaped in json file
+        self.assertTrue('黄' in json_db['6686-1']['description'])
+        self.assertFalse('黄' in json_text.decode('utf-8'))
+        self.assertTrue(b'\\u9ec4' in json_text)
+
+        # Make sure the id has no prefix
+        self.assertEqual(json_db['6686-1']['id'], '6686-1')
+
+    def test_convert_pickle_prefix(self):
+        '''test converting pickle file to json with prefix'''
+        pickle_file = 'scripts/testfiles/database-protv2.pickle'
+        json_file = os.path.join(self.tmpdir, 'test.json')
+
+        rc = subprocess.call(["scripts/convert-pickle.py",
+                              "-i", pickle_file,
+                              "-o", json_file,
+                              "-p", "TEST-"])
+        self.assertEqual(rc, 0)
+
+        json_db = self._load_json(json_file)
+
+        self.assertEqual(json_db['6686-1']['id'], 'TEST-6686-1')
+
+    def test_convert_pickle_v3(self):
+        '''test converting pickle v3 file to json'''
+        pickle_file = 'scripts/testfiles/database-protv3.pickle'
+        json_file = os.path.join(self.tmpdir, 'test.json')
+
+        rc = subprocess.call(["scripts/convert-pickle.py",
+                              "-i", pickle_file,
+                              "-o", json_file])
+        self.assertEqual(rc, 0)
+
+        pickle_db = self._load_pickle(pickle_file)
+        json_db = self._load_json(json_file)
+        json_text = self._load_text(json_file)
+
+        # Make sure the databases are the same
+        self.assertEqual(pickle_db, json_db)
+
+        # Make sure unicode is escaped in json file
+        self.assertTrue('黄' in json_db['6686-1']['description'])
+        self.assertFalse('黄' in json_text.decode('utf-8'))
+        self.assertTrue(b'\\u9ec4' in json_text)
+
+        # Make sure the id has no prefix
+        self.assertEqual(json_db['6686-1']['id'], '6686-1')
+
diff --git a/scripts/testfiles/database-protv2.pickle b/scripts/testfiles/database-protv2.pickle
 new file mode 100644
index 0000000..285e2ef
 Binary files /dev/null and b/scripts/testfiles/database-protv2.pickle differ
diff --git a/scripts/testfiles/database-protv3.pickle b/scripts/testfiles/database-protv3.pickle
 new file mode 100644
index 0000000..2c94bcd
 Binary files /dev/null and b/scripts/testfiles/database-protv3.pickle differ