From c401f2c6cc5f0aa4c6de245dcf252d3bbcc920fd Mon Sep 17 00:00:00 2001
From: J08nY
Date: Thu, 27 Jul 2017 17:30:51 +0200
Subject: Add get_signature and get_encrypted to wrappers.

---
 src/mailman_pgp/pgp/inline.py  | 22 ++++++++++++++++++++++
 src/mailman_pgp/pgp/mime.py    | 18 ++++++++++++++++++
 src/mailman_pgp/pgp/wrapper.py | 21 +++++++++++++++++++++
 3 files changed, 61 insertions(+)

diff --git a/src/mailman_pgp/pgp/inline.py b/src/mailman_pgp/pgp/inline.py
index d505585..d2ea05e 100644
--- a/src/mailman_pgp/pgp/inline.py
+++ b/src/mailman_pgp/pgp/inline.py
@@ -75,6 +75,10 @@ class InlineWrapper:
         return any(self._walk(self._is_signed))
 
     def get_signed(self):
+        """
+
+        :return:
+        """
         for part in walk(self.msg):
             if not part.is_multipart() and self._is_signed(part):
                 try:
@@ -83,6 +87,15 @@ class InlineWrapper:
                     continue
                 yield msg
 
+    def get_signature(self):
+        """
+
+        :return:
+        """
+        for part in walk(self.msg):
+            if not part.is_multipart() and self._is_signed(part):
+                yield part.get_payload()
+
     def _is_encrypted(self, part):
         try:
             msg = PGPMessage.from_blob(part.get_payload())
@@ -109,6 +122,15 @@ class InlineWrapper:
         """
         return any(self._walk(self._is_encrypted))
 
+    def get_encrypted(self):
+        """
+
+        :return:
+        """
+        for part in walk(self.msg):
+            if not part.is_multipart() and self._is_encrypted(part):
+                yield part.get_payload()
+
     def _has_keys(self, part):
         try:
             dearm = Armorable.ascii_unarmor(part.get_payload())
diff --git a/src/mailman_pgp/pgp/mime.py b/src/mailman_pgp/pgp/mime.py
index 9978b54..c702eb8 100644
--- a/src/mailman_pgp/pgp/mime.py
+++ b/src/mailman_pgp/pgp/mime.py
@@ -88,8 +88,19 @@ class MIMEWrapper:
         return self.is_signed()
 
     def get_signed(self):
+        """
+
+        :return:
+        """
         yield self.msg.get_payload(0).as_string()
 
+    def get_signature(self):
+        """
+
+        :return:
+        """
+        yield self.msg.get_payload(1).get_payload()
+
     def is_encrypted(self):
         """
         Whether the whole message is MIME encrypted as per RFC3156 section 4.
@@ -115,6 +126,13 @@ class MIMEWrapper:
     def has_encryption(self):
         return self.is_encrypted()
 
+    def get_encrypted(self):
+        """
+
+        :return:
+        """
+        yield self.msg.get_payload(1).get_payload()
+
     def is_keys(self):
         """
         Whether the message has only keys as per RFC3156 section 7.
diff --git a/src/mailman_pgp/pgp/wrapper.py b/src/mailman_pgp/pgp/wrapper.py
index cbfa7f2..a1041ad 100644
--- a/src/mailman_pgp/pgp/wrapper.py
+++ b/src/mailman_pgp/pgp/wrapper.py
@@ -81,6 +81,16 @@ class PGPWrapper():
         elif self.inline.is_signed():
             yield from self.inline.get_signed()
 
+    def get_signature(self):
+        """
+
+        :return:
+        """
+        if self.mime.is_signed():
+            yield from self.mime.get_signature()
+        elif self.inline.is_signed():
+            yield from self.inline.get_signature()
+
     def sign(self, key, **kwargs):
         """
         Sign a message with key.
@@ -129,6 +139,17 @@ class PGPWrapper():
         """
         return self.mime.has_encryption() or self.inline.has_encryption()
 
+
+    def get_encrypted(self):
+        """
+
+        :return:
+        """
+        if self.mime.is_signed():
+            yield from self.mime.get_encrypted()
+        elif self.inline.is_signed():
+            yield from self.inline.get_encrypted()
+
     def encrypt(self, *keys, **kwargs):
         """
         Encrypt the message with key/s, using cipher.
-- 
cgit v1.2.3-70-g09d2