readPemFromFile() improved to select one of possible many PEM headers

1 files changed, 11 insertions, 15 deletions

diff --git a/tools/pkcs8dump.py b/tools/pkcs8dump.py
index c3e4dc6..ed13875 100644
--- a/tools/pkcs8dump.py
+++ b/tools/pkcs8dump.py
@@ -1,7 +1,6 @@
 # Read  bunch of ASN.1/PEM plain/encrypted private keys in PKCS#8 
 # format on stdin, parse each into plain text, then build substrate from it
 from pyasn1.codec.der import decoder, encoder
-from pyasn1 import error
 from pyasn1_modules import rfc5208, pem
 import sys
 
@@ -10,26 +9,23 @@ if len(sys.argv) != 1:
 $ cat pkcs8key.pem | %s""" % (sys.argv[0], sys.argv[0]))
     sys.exit(-1)
     
-keyTypePlain = rfc5208.PrivateKeyInfo()
-keyTypeEncrypted = rfc5208.EncryptedPrivateKeyInfo()
-
 cnt = 0
 
 while 1:
-    substrate = pem.readPemFromFile(
-                    sys.stdin,
-                    ('-----BEGIN PRIVATE KEY-----',
-                     '-----BEGIN ENCRYPTED PRIVATE KEY-----'),
-                    ('-----END PRIVATE KEY-----',
-                     '-----END ENCRYPTED PRIVATE KEY-----')
-                )
+    idx, substrate = pem.readPemFromFile(sys.stdin, ('-----BEGIN PRIVATE KEY-----', '-----END PRIVATE KEY-----'), ('-----BEGIN ENCRYPTED PRIVATE KEY-----', '-----END ENCRYPTED PRIVATE KEY-----') )
     if not substrate:
         break
 
-    try:        
-        key, rest = decoder.decode(substrate, asn1Spec=keyTypeEncrypted)
-    except error.PyAsn1Error:
-        key, rest = decoder.decode(substrate, asn1Spec=keyTypePlain)
+    if idx == 0:
+        key, rest = decoder.decode(
+                       substrate, asn1Spec=rfc5208.PrivateKeyInfo()
+                    )
+    elif idx == 1:
+        key, rest = decoder.decode(
+                       substrate, asn1Spec=rfc5208.EncryptedPrivateKeyInfo()
+                    )
+    else:
+        break
 
     if rest: substrate = substrate[:-len(rest)]