[feat] use a config class for the protect function

5 files changed, 46 insertions, 27 deletions

diff --git a/memoryhole/__init__.py b/memoryhole/__init__.py
index fd621ea..a7b0df5 100644
--- a/memoryhole/__init__.py
+++ b/memoryhole/__init__.py
@@ -1,12 +1,8 @@
-from memoryhole.protection import protect
-from memoryhole.openpgp import OpenPGP
+from memoryhole.protection import protect, ProtectConfig
+from memoryhole.openpgp import IOpenPGP
 from memoryhole.gpg import Gnupg
 
 
-PROTECTED_HEADERS = ('Subject', 'Message-ID', 'Date', 'To', 'From')
-OBSCURED_HEADERS = ('Subject', 'Message-ID', 'Date', 'To', 'From')
-
-
 def unwrap(msg, opengp=Gnupg()):
     """
     Unwrap an email replacing and verifying memory hole headers.
@@ -17,10 +13,10 @@ def unwrap(msg, opengp=Gnupg()):
                     verification
     :type openpgp: OpenPGP
 
-    :return: a dencrypted email
+    :return: a decrypted email
     :rtype: Message
     """
     raise NotImplementedError()
 
 
-__all__ = ["protect", "unwrap", "OpenPGP"]
+__all__ = ["protect", "ProtectConfig", "unwrap", "IOpenPGP"]
diff --git a/memoryhole/gpg.py b/memoryhole/gpg.py
index 533d12a..3f1c59a 100644
--- a/memoryhole/gpg.py
+++ b/memoryhole/gpg.py
@@ -1,9 +1,9 @@
 from zope.interface import implementer
 
-from memoryhole.openpgp import OpenPGP
+from memoryhole.openpgp import IOpenPGP
 
 
-@implementer(OpenPGP)
+@implementer(IOpenPGP)
 class Gnupg(object):
     def __init__(self):
         from gnupg import GPG
@@ -11,7 +11,7 @@ class Gnupg(object):
 
     def encrypt(self, data, encraddr, signaddr):
         result = self.gpg.encrypt(data, *encraddr, default_key=signaddr)
-        self._assert_gpg_result_ok(result)
+        self._check_gpg_error(result)
         return result.data
 
     def decrypt(self, data):
@@ -20,7 +20,7 @@ class Gnupg(object):
     def verify(self, data, signature):
         pass
 
-    def _assert_gpg_result_ok(self, result):
+    def _check_gpg_error(self, result):
         stderr = getattr(result, 'stderr', '')
         if getattr(result, 'ok', False) is not True:
             raise RuntimeError('Failed to encrypt/decrypt: %s' % stderr)
diff --git a/memoryhole/openpgp.py b/memoryhole/openpgp.py
index ac5cb6e..3cefc3c 100644
--- a/memoryhole/openpgp.py
+++ b/memoryhole/openpgp.py
@@ -1,7 +1,7 @@
 from zope.interface import Interface
 
 
-class OpenPGP(Interface):
+class IOpenPGP(Interface):
     def encrypt(data, encraddr, singaddr):
         """
         Encrypt and sign data.
diff --git a/memoryhole/protection.py b/memoryhole/protection.py
index b064f47..d1b372b 100644
--- a/memoryhole/protection.py
+++ b/memoryhole/protection.py
@@ -5,31 +5,52 @@ from memoryhole.gpg import Gnupg
 from memoryhole.rfc3156 import PGPEncrypted, MultipartEncrypted
 
 
-def protect(msg, openpgp=Gnupg(), encrypt=True, obscure=True):
+class ProtectConfig(object):
+
+    PROTECTED_HEADERS = ('Subject', 'Message-ID', 'Date', 'To', 'From')
+    OBSCURED_HEADERS = ('Subject', 'Message-ID', 'Date', 'To', 'From')
+
+    def __init__(self, openpgp=None, protected_headers=PROTECTED_HEADERS,
+                 obscured_headers=OBSCURED_HEADERS):
+        """
+        Configuration parameters for the protection
+
+        :param openpgp: the implementation of openpgp to use for encryption
+                        and/or signature
+        :type openpgp: IOpenPGP
+        :param protected_headers: list of headers to protect
+        :type protected_headers: [str]
+        :param obscured_headers: list of headers to obscure
+        :type obscured_headers: [str]
+        """
+        if openpgp is None:
+            openpgp = Gnupg()
+        self.openpgp = openpgp
+
+
+def protect(msg, encrypt=True, config=None):
     """
     Protect an email with memory hole. It will protect the PROTECTED_HEADERS
     and if obscure=True will obscure the OBSCURED_HEADERS
 
     :param msg: the email to be protected
     :type msg: Message
-    :param openpgp: the implementation of openpgp to use for encryption and/or
-                    signature
-    :type openpgp: OpenPGP
     :param encrypt: should the message be encrypted
     :type encrypt: bool
-    :param obscure: should the headers be obscured
-    :type obsucre: bool
 
     :return: an encrypted and/or signed email
     :rtype: Message
     """
+    if config is None:
+        config = ProtectConfig()
+
     if encrypt:
-        return _encrypt_mime(msg, openpgp)
+        return _encrypt_mime(msg, config)
 
     raise NotImplementedError()
 
 
-def _encrypt_mime(msg, openpgp):
+def _encrypt_mime(msg, config):
     encraddr = _recipient_addresses(msg)
     signaddr = _from_address(msg)
 
@@ -38,8 +59,8 @@ def _encrypt_mime(msg, openpgp):
         newmsg.add_header(hkey, hval)
         del(msg[hkey])
 
-    encstr = openpgp.encrypt(msg.as_string(unixfrom=False),
-                             encraddr, signaddr)
+    encstr = config.openpgp.encrypt(msg.as_string(unixfrom=False),
+                                    encraddr, signaddr)
     encmsg = MIMEApplication(
         encstr, _subtype='octet-stream', _encoder=lambda x: x)
     encmsg.add_header('content-disposition', 'attachment',
diff --git a/tests/test_protection.py b/tests/test_protection.py
index 59effaa..1220639 100644
--- a/tests/test_protection.py
+++ b/tests/test_protection.py
@@ -2,7 +2,7 @@ import unittest
 from email.parser import Parser
 from zope.interface import implementer
 
-from memoryhole import protect, OpenPGP
+from memoryhole import protect, ProtectConfig, IOpenPGP
 
 
 FROM = "me@domain.com"
@@ -27,7 +27,8 @@ class ProtectTest(unittest.TestCase):
         p = Parser()
         msg = p.parsestr(EMAIL)
         encrypter = Encrypter()
-        encmsg = protect(msg, encrypter)
+        conf = ProtectConfig(openpgp=encrypter)
+        encmsg = protect(msg, config=conf)
 
         self.assertEqual(encmsg.get_payload(1).get_payload(), encrypter.encstr)
         self.assertEqual(BODY, encrypter.data[1:-1])  # remove '\n'
@@ -39,14 +40,15 @@ class ProtectTest(unittest.TestCase):
         p = Parser()
         msg = p.parsestr(EMAIL)
         encrypter = Encrypter()
-        encmsg = protect(msg, encrypter, obscure=False)
+        conf = ProtectConfig(openpgp=encrypter, obscured_headers=[])
+        encmsg = protect(msg, config=conf)
 
         self.assertEqual(encmsg['from'], FROM)
         self.assertEqual(encmsg['to'], TO)
         self.assertEqual(encmsg['subject'], SUBJECT)
 
 
-@implementer(OpenPGP)
+@implementer(IOpenPGP)
 class Encrypter(object):
     encstr = "this is encrypted"