Repository: jaywink/federation

Issue:
Add support for receiving new style Diaspora private message payloads
See https://github.com/diaspora/diaspora_federation/issues/30

Discussion:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 63c2f0f..3ba3370 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## [unreleased]
 
+### Added
+* New style Diaspora private encrypted JSON payloads are now supported in the receiving side. Outbound private Diaspora payloads are still sent as legacy encrypted payloads. ([issue](https://github.com/jaywink/federation/issues/83))
+    * No additional changes need to be made when calling `handle_receive` from your task processing. Just pass in the full received XML or JSON payload as a string with recipient user object as before.
+
 ### Fixed
 * Fix getting sender from a combination of legacy Diaspora encrypted payload and new entity names (for example `author`). This combination probably only existed in this library.
 * Correctly extend entity `_children`. Certain Diaspora payloads caused `_children` for an entity to be written over by an empty list, causing for example status message photos to not be saved. Correctly do an extend on it. ([issue](https://github.com/jaywink/federation/issues/89))
@@ -9,6 +13,7 @@
 
 ### Removed
 * `Post.photos` entity attribute was never used by any code and has been removed. Child entities of type `Image` are stored in the `Post._children` as before.
+* Removed deprecated user private key lookup using `user.key` in Diaspora receive processing. Passed in `user` objects must now have a `private_key` attribute. 
 
 ## [0.12.0] - 2017-05-22
 
diff --git a/federation/protocols/base.py b/federation/protocols/base.py
index 0f48de3..e298ce0 100644
--- a/federation/protocols/base.py
+++ b/federation/protocols/base.py
@@ -37,7 +37,7 @@ class BaseProtocol(object):
         """
         raise NotImplementedError("Implement in subclass")
 
-    def receive(self, payload, user=None, sender_key_fetcher=None, *args, **kwargs):
+    def receive(self, payload, user=None, sender_key_fetcher=None):
         """Receive a payload.
 
         Args:
diff --git a/federation/protocols/diaspora/encrypted.py b/federation/protocols/diaspora/encrypted.py
new file mode 100644
index 0000000..ccf23a4
--- /dev/null
+++ b/federation/protocols/diaspora/encrypted.py
@@ -0,0 +1,30 @@
+import json
+from base64 import b64decode
+
+from Crypto.Cipher import PKCS1_v1_5, AES
+from lxml import etree
+
+
+def pkcs7_unpad(data):
+    """Remove the padding bytes that were added at point of encryption."""
+    if isinstance(data, str):
+        return data[0:-ord(data[-1])]
+    else:
+        return data[0:-data[-1]]
+
+
+class EncryptedPayload:
+    """Diaspora encrypted JSON payloads."""
+
+    @staticmethod
+    def decrypt(payload, private_key):
+        """Decrypt an encrypted JSON payload and return the Magic Envelope document inside."""
+        cipher = PKCS1_v1_5.new(private_key)
+        aes_key_str = cipher.decrypt(b64decode(payload.get("aes_key")), sentinel=None)
+        aes_key = json.loads(aes_key_str.decode("utf-8"))
+        key = b64decode(aes_key.get("key"))
+        iv = b64decode(aes_key.get("iv"))
+        encrypted_magic_envelope = b64decode(payload.get("encrypted_magic_envelope"))
+        encrypter = AES.new(key, AES.MODE_CBC, iv)
+        content = encrypter.decrypt(encrypted_magic_envelope)
+        return etree.fromstring(pkcs7_unpad(content))
diff --git a/federation/protocols/diaspora/magic_envelope.py b/federation/protocols/diaspora/magic_envelope.py
index ff9cc13..46a9a08 100644
--- a/federation/protocols/diaspora/magic_envelope.py
+++ b/federation/protocols/diaspora/magic_envelope.py
@@ -8,7 +8,7 @@ from lxml import etree
 NAMESPACE = "http://salmon-protocol.org/ns/magic-env"
 
 
-class MagicEnvelope():
+class MagicEnvelope:
     """Diaspora protocol magic envelope.
 
     See: http://diaspora.github.io/diaspora_federation/federation/magicsig.html
diff --git a/federation/protocols/diaspora/protocol.py b/federation/protocols/diaspora/protocol.py
index e65d3ef..73c417e 100644
--- a/federation/protocols/diaspora/protocol.py
+++ b/federation/protocols/diaspora/protocol.py
@@ -1,6 +1,5 @@
 import json
 import logging
-import warnings
 from base64 import b64decode, urlsafe_b64decode, b64encode, urlsafe_b64encode
 from urllib.parse import unquote_plus
 
@@ -13,7 +12,9 @@ from lxml import etree
 
 from federation.exceptions import EncryptedMessageError, NoSenderKeyFoundError, SignatureVerificationError
 from federation.protocols.base import BaseProtocol
+from federation.protocols.diaspora.encrypted import EncryptedPayload
 from federation.protocols.diaspora.magic_envelope import MagicEnvelope
+from federation.utils.text import decode_if_bytes
 
 logger = logging.getLogger("federation")
 
@@ -29,7 +30,7 @@ def identify_payload(payload):
     """
     # Private encrypted JSON payload
     try:
-        data = json.loads(payload)
+        data = json.loads(decode_if_bytes(payload))
         if "encrypted_magic_envelope" in data:
             return True
     except Exception:
@@ -55,17 +56,36 @@ class Protocol(BaseProtocol):
 
     Mostly taken from Pyaspora (https://github.com/lukeross/pyaspora).
     """
-    def receive(self, payload, user=None, sender_key_fetcher=None, skip_author_verification=False, *args, **kwargs):
+    def __init__(self):
+        super().__init__()
+        self.encrypted = self.legacy = False
+
+    def get_json_payload_magic_envelope(self, payload):
+        """Encrypted JSON payload"""
+        private_key = self._get_user_key(self.user)
+        return EncryptedPayload.decrypt(payload=payload, private_key=private_key)
+
+    def store_magic_envelope_doc(self, payload):
+        """Get the Magic Envelope, trying JSON first."""
+        try:
+            json_payload = json.loads(decode_if_bytes(payload))
+        except ValueError:
+            # XML payload
+            xml = unquote_plus(payload)
+            xml = xml.lstrip().encode("utf-8")
+            logger.debug("diaspora.protocol.store_magic_envelope_doc: xml payload: %s", xml)
+            self.doc = etree.fromstring(xml)
+        else:
+            logger.debug("diaspora.protocol.store_magic_envelope_doc: json payload: %s", json_payload)
+            self.doc = self.get_json_payload_magic_envelope(json_payload)
+
+    def receive(self, payload, user=None, sender_key_fetcher=None, skip_author_verification=False):
         """Receive a payload.
 
         For testing purposes, `skip_author_verification` can be passed. Authorship will not be verified."""
         self.user = user
         self.get_contact_key = sender_key_fetcher
-        # Prepare payload
-        xml = unquote_plus(payload)
-        xml = xml.lstrip().encode("utf-8")
-        logger.debug("diaspora.protocol.receive: xml content: %s", xml)
-        self.doc = etree.fromstring(xml)
+        self.store_magic_envelope_doc(payload)
         # Check for a legacy header
         self.find_header()
         # Open payload and get actual message
@@ -78,17 +98,11 @@ class Protocol(BaseProtocol):
         return self.sender_handle, self.content
 
     def _get_user_key(self, user):
-        if not hasattr(self.user, "private_key") or not self.user.private_key:
-            if hasattr(self.user, "key") and self.user.key:
-                warnings.warn("Using `key` in user object for private key has been deprecated. Please "
-                              "have available `private_key` instead. Usage of `key` will be removed after 0.8.0.",
-                              DeprecationWarning)
-                return self.user.key
+        if not getattr(self.user, "private_key", None):
             raise EncryptedMessageError("Cannot decrypt private message without user key")
         return self.user.private_key
 
     def find_header(self):
-        self.encrypted = self.legacy = False
         self.header = self.doc.find(".//{"+PROTOCOL_NS+"}header")
         if self.header != None:
             # Legacy public header found
diff --git a/federation/utils/text.py b/federation/utils/text.py
new file mode 100644
index 0000000..310a5e1
--- /dev/null
+++ b/federation/utils/text.py
@@ -0,0 +1,5 @@
+def decode_if_bytes(text):
+    try:
+        return text.decode("utf-8")
+    except AttributeError:
+        return text


Test patch:
diff --git a/federation/tests/protocols/diaspora/test_encrypted.py b/federation/tests/protocols/diaspora/test_encrypted.py
new file mode 100644
index 0000000..5d10988
--- /dev/null
+++ b/federation/tests/protocols/diaspora/test_encrypted.py
@@ -0,0 +1,33 @@
+from unittest.mock import patch, Mock
+
+from Crypto.Cipher import AES
+
+from federation.protocols.diaspora.encrypted import pkcs7_unpad, EncryptedPayload
+
+
+def test_pkcs7_unpad():
+    assert pkcs7_unpad(b"foobar\x02\x02") == b"foobar"
+    assert pkcs7_unpad("foobar\x02\x02") == "foobar"
+
+
+class TestEncryptedPayload:
+    @patch("federation.protocols.diaspora.encrypted.PKCS1_v1_5.new")
+    @patch("federation.protocols.diaspora.encrypted.AES.new")
+    @patch("federation.protocols.diaspora.encrypted.pkcs7_unpad", side_effect=lambda x: x)
+    @patch("federation.protocols.diaspora.encrypted.b64decode", side_effect=lambda x: x)
+    def test_decrypt(self, mock_decode, mock_unpad, mock_aes, mock_pkcs1):
+        mock_decrypt = Mock(return_value=b'{"iv": "foo", "key": "bar"}')
+        mock_pkcs1.return_value = Mock(decrypt=mock_decrypt)
+        mock_encrypter = Mock(return_value="<foo>bar</foo>")
+        mock_aes.return_value = Mock(decrypt=mock_encrypter)
+        doc = EncryptedPayload.decrypt(
+            {"aes_key": '{"iv": "foo", "key": "bar"}', "encrypted_magic_envelope": "magically encrypted"},
+            "private_key",
+        )
+        mock_pkcs1.assert_called_once_with("private_key")
+        mock_decrypt.assert_called_once_with('{"iv": "foo", "key": "bar"}', sentinel=None)
+        assert mock_decode.call_count == 4
+        mock_aes.assert_called_once_with("bar", AES.MODE_CBC, "foo")
+        mock_encrypter.assert_called_once_with("magically encrypted")
+        assert doc.tag == "foo"
+        assert doc.text == "bar"
diff --git a/federation/tests/protocols/diaspora/test_protocol.py b/federation/tests/protocols/diaspora/test_protocol.py
index 3fd5a37..f33a31a 100644
--- a/federation/tests/protocols/diaspora/test_protocol.py
+++ b/federation/tests/protocols/diaspora/test_protocol.py
@@ -155,3 +155,24 @@ class TestDiasporaProtocol(DiasporaTestBase):
         data = protocol.build_send(entity, from_user)
         mock_init_message.assert_called_once_with(mock_entity_xml, from_user.handle, from_user.private_key)
         assert data == {"xml": "xmldata"}
+
+    @patch("federation.protocols.diaspora.protocol.EncryptedPayload.decrypt")
+    def test_get_json_payload_magic_envelope(self, mock_decrypt):
+        protocol = Protocol()
+        protocol.user = MockUser()
+        protocol.get_json_payload_magic_envelope("payload")
+        mock_decrypt.assert_called_once_with(payload="payload", private_key="foobar")
+
+    @patch.object(Protocol, "get_json_payload_magic_envelope", return_value=etree.fromstring("<foo>bar</foo>"))
+    def test_store_magic_envelope_doc_json_payload(self, mock_store):
+        protocol = Protocol()
+        protocol.store_magic_envelope_doc('{"foo": "bar"}')
+        mock_store.assert_called_once_with({"foo": "bar"})
+        assert protocol.doc.tag == "foo"
+        assert protocol.doc.text == "bar"
+
+    def test_store_magic_envelope_doc_xml_payload(self):
+        protocol = Protocol()
+        protocol.store_magic_envelope_doc("<foo>bar</foo>")
+        assert protocol.doc.tag == "foo"
+        assert protocol.doc.text == "bar"
diff --git a/federation/tests/utils/test_text.py b/federation/tests/utils/test_text.py
new file mode 100644
index 0000000..7f629d3
--- /dev/null
+++ b/federation/tests/utils/test_text.py
@@ -0,0 +1,6 @@
+from federation.utils.text import decode_if_bytes
+
+
+def test_decode_if_bytes():
+    assert decode_if_bytes(b"foobar") == "foobar"
+    assert decode_if_bytes("foobar") == "foobar"