Merge pull request #375 from tlsfuzzer/shadchin-patch-1

Fix tests with new Python

Author: tomato42

src/ecdsa/der.py

@@ -464,13 +464,8 @@ def unpem(pem):
     if isinstance(pem, text_type):  # pragma: no branch
         pem = pem.encode()
 
-    d = b"".join(
-        [
-            l.strip()
-            for l in pem.split(b"\n")
-            if l and not l.startswith(b"-----")
-        ]
-    )
+    lines = (l.strip() for l in pem.split(b"\n"))
+    d = b"".join(l for l in lines if l and not l.startswith(b"-----"))
     return base64.b64decode(d)
 
 

src/ecdsa/test_der.py

@@ -26,6 +26,8 @@
     remove_octet_string,
     remove_sequence,
     encode_implicit,
+    unpem,
+    topem,
 )
 
 
@@ -629,3 +631,16 @@ def test_remove_implicit_rejects_truncated_length():
         UnexpectedDER, match="Length longer than the provided buffer"
     ):
         remove_implicit(bad)
+
+
+def test_unpem_with_str():
+    data = b"\xff\x00\x01"
+    encoded = topem(data, "KEY")
+    if sys.version_info < (3, 0):
+        # on python2 the string can be unicode
+        encoded = unicode(encoded)
+    else:
+        # on python3 topem() returns bytes
+        encoded = str(encoded, "ASCII")
+    out = unpem(encoded)
+    assert data == out