diff --git a/src/ecdsa/der.py b/src/ecdsa/der.py index 5d35d698..fc31a50e 100644 --- a/src/ecdsa/der.py +++ b/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) diff --git a/src/ecdsa/test_der.py b/src/ecdsa/test_der.py index 26909bc6..bbddf5eb 100644 --- a/src/ecdsa/test_der.py +++ b/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