From 8409ac1fc0625c1c7c08a13b457408077fa47095 Mon Sep 17 00:00:00 2001
From: Hynek Schlawack <hs@ox.cx>
Date: Tue, 20 Jun 2023 11:20:25 +0200
Subject: [PATCH] Add support for headers

Ref #45
---
 CHANGELOG.md       |  1 +
 docs/api.rst       |  2 +-
 src/pem/_core.py   | 22 ++++++++++++++++++++++
 tests/data.py      | 21 +++++++++++++++++++--
 tests/test_core.py | 39 ++++++++++++++++++++++++++++++++++++++-
 5 files changed, 81 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 425c47b..26a2af2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -21,6 +21,7 @@ You can find out backwards-compatibility policy [here](https://github.com/hynek/
 
 - Support for RFC 4880 OpenPGP private & public keys: `pem.OpenPGPPublicKey` and `pem.OpenPGPPrivateKey`.
   [#72](https://github.com/hynek/pem/issues/72)
+- Support for intra-payload headers like the ones used in OpenPGP keys using the `meta_headers` property.
 - `pem.parse_file()` now accepts also [`pathlib.Path`](https://docs.python.org/3/library/pathlib.html#pathlib.Path) objects.
 - Added `payload_as_text()`, `payload_as_bytes()` and `payload_decoded()` to all PEM objects that allow to directly access the payload without the envelope and possible headers.
   [#74](https://github.com/hynek/pem/pull/74)
diff --git a/docs/api.rst b/docs/api.rst
index 8438eda..4360447 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -47,7 +47,7 @@ The following objects can be returned by the parsing functions.
 Their shared provided API is minimal:
 
 .. autoclass:: AbstractPEMObject
-   :members: __str__, as_bytes, as_text, sha1_hexdigest, payload_as_bytes, payload_as_text, payload_decoded
+   :members: __str__, as_bytes, as_text, sha1_hexdigest, payload_as_bytes, payload_as_text, payload_decoded, meta_headers
 
 
 Twisted
diff --git a/src/pem/_core.py b/src/pem/_core.py
index e414ed9..1f3bb8b 100644
--- a/src/pem/_core.py
+++ b/src/pem/_core.py
@@ -120,6 +120,28 @@ def payload_decoded(self) -> bytes:
         """
         return b64decode(self._extracted_payload)
 
+    @cached_property
+    def meta_headers(self) -> dict[str, str]:
+        """
+        Return a dictionary of payload headers.
+
+        .. versionadded:: 23.1.0
+        """
+        expl = {}
+        for line in self._pem_bytes.decode().splitlines()[1:-1]:
+            if ":" not in line:
+                break
+
+            key, val = line.split(": ", 1)
+
+            # Strip quotes if they're only at the beginning and end.
+            if val.count('"') == 2 and val[0] == '"' and val[-1] == '"':
+                val = val[1:-1]
+
+            expl[key] = val
+
+        return expl
+
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, type(self)):
             return NotImplemented
diff --git a/tests/data.py b/tests/data.py
index 2e9619e..678e6c4 100644
--- a/tests/data.py
+++ b/tests/data.py
@@ -308,8 +308,7 @@
 # Taken from https://tools.ietf.org/html/rfc4716#section-3.6.
 KEY_PEM_RFC4716_PUBLIC = rb"""---- BEGIN SSH2 PUBLIC KEY ----
 Subject: me
-Comment: 1024-bit rsa, created by me@example.com Mon Jan 15 \
-08:31:24 2001
+Comment: 1024-bit rsa, created by me@example.com Mon Jan 15 08:31:24 2001
 AAAAB3NzaC1yc2EAAAABJQAAAIEAiPWx6WM4lhHNedGfBpPJNPpZ7yKu+dnn1SJejgt4
 596k6YjzGGphH2TUxwKzxcKDKKezwkpfnxPkSMkuEspGRt/aZZ9wa++Oi7Qkr8prgHc4
 soW6NUlfDzpvZK2H5E7eQaSeP3SAwGmQKUFHCddNaP0L+hM7zhFNzjFvpaMgJw0=
@@ -436,3 +435,21 @@
 =n8OM
 -----END PGP PRIVATE KEY BLOCK-----
 """
+
+# From https://datatracker.ietf.org/doc/html/rfc7468#section-5.2
+CERT_PEM_CERTIFICATE_EXPLANATORY = b"""\
+Subject: CN=Atlantis
+Issuer: CN=Atlantis
+Validity: from 7/9/2012 3:10:38 AM UTC to 7/9/2013 3:10:37 AM UTC
+-----BEGIN CERTIFICATE-----
+MIIBmTCCAUegAwIBAgIBKjAJBgUrDgMCHQUAMBMxETAPBgNVBAMTCEF0bGFudGlz
+MB4XDTEyMDcwOTAzMTAzOFoXDTEzMDcwOTAzMTAzN1owEzERMA8GA1UEAxMIQXRs
+YW50aXMwXDANBgkqhkiG9w0BAQEFAANLADBIAkEAu+BXo+miabDIHHx+yquqzqNh
+Ryn/XtkJIIHVcYtHvIX+S1x5ErgMoHehycpoxbErZmVR4GCq1S2diNmRFZCRtQID
+AQABo4GJMIGGMAwGA1UdEwEB/wQCMAAwIAYDVR0EAQH/BBYwFDAOMAwGCisGAQQB
+gjcCARUDAgeAMB0GA1UdJQQWMBQGCCsGAQUFBwMCBggrBgEFBQcDAzA1BgNVHQEE
+LjAsgBA0jOnSSuIHYmnVryHAdywMoRUwEzERMA8GA1UEAxMIQXRsYW50aXOCASow
+CQYFKw4DAh0FAANBAKi6HRBaNEL5R0n56nvfclQNaXiDT174uf+lojzA4lhVInc0
+ILwpnZ1izL4MlI9eCSHhVQBHEp2uQdXJB+d5Byg=
+-----END CERTIFICATE-----
+"""
diff --git a/tests/test_core.py b/tests/test_core.py
index 9adc206..c4904a6 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -112,7 +112,7 @@ def test_payload_as_text(self):
             (KEY_PEM_SSHCOM_PRIVATE, b"Comment:"),
             (KEY_PEM_OPENPGP_PUBLIC, b"Version:"),
             (KEY_PEM_OPENPGP_PRIVATE, b"Comment:"),
-            (KEY_PEM_PKCS5_ENCRYPTED, (b"Proc-Type:")),
+            (KEY_PEM_PKCS5_ENCRYPTED, b"Proc-Type:"),
         ],
     )
     def test_payload_with_headers(self, bs, forbidden):
@@ -663,3 +663,40 @@ def test_openpgp_private_key(self):
 
         assert isinstance(key, pem.OpenPGPPrivateKey)
         assert KEY_PEM_OPENPGP_PRIVATE == key.as_bytes()
+
+    @pytest.mark.parametrize(
+        "bs, hdrs",
+        [
+            (KEY_PEM_SSHCOM_PRIVATE, {"Comment": "rsa-key-20210120"}),
+            (
+                KEY_PEM_OPENPGP_PUBLIC,
+                {"Version": "Encryption Desktop 10.4.2 (Build 289)"},
+            ),
+            (
+                KEY_PEM_OPENPGP_PRIVATE,
+                {
+                    "Comment": "https://www.ietf.org/id/draft-bre-openpgp-samples-01.html"
+                },
+            ),
+            (
+                KEY_PEM_PKCS5_ENCRYPTED,
+                {
+                    "DEK-Info": "DES-EDE3-CBC,8A72BD2DC1C9092F",
+                    "Proc-Type": "4,ENCRYPTED",
+                },
+            ),
+        ],
+    )
+    def test_headers(self, bs, hdrs):
+        """
+        Headers are preserved.
+        """
+        assert hdrs == pem.parse(bs)[0].meta_headers
+
+    def test_no_headers(self):
+        """
+        No headers, no problem.
+        """
+        cert = pem.parse(CERT_PEM_OPENSSL_TRUSTED)[0]
+
+        assert {} == cert.meta_headers