From 64b1c8c02e753fa55084f0a1065d1fa8e75d5b5e Mon Sep 17 00:00:00 2001
From: Jason Robinson <mail@jasonrobinson.me>
Date: Tue, 4 Jul 2017 23:14:02 +0300
Subject: [PATCH 1/7] Minor code style cleanup

---
 federation/protocols/base.py                    | 2 +-
 federation/protocols/diaspora/magic_envelope.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

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/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

From ccf161a5d3df60b2d46d9363d1ea138641b2a32e Mon Sep 17 00:00:00 2001
From: Jason Robinson <mail@jasonrobinson.me>
Date: Tue, 4 Jul 2017 23:14:42 +0300
Subject: [PATCH 2/7] Remove deprecated 'user.key' lookup for private key

---
 CHANGELOG.md                              | 1 +
 federation/protocols/diaspora/protocol.py | 7 +------
 2 files changed, 2 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 63c2f0f..904a239 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,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/diaspora/protocol.py b/federation/protocols/diaspora/protocol.py
index e65d3ef..c0eb9df 100644
--- a/federation/protocols/diaspora/protocol.py
+++ b/federation/protocols/diaspora/protocol.py
@@ -78,12 +78,7 @@ 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
 

From 2e8d6082569ca7c1e4bab26f4ac519b44ee5624e Mon Sep 17 00:00:00 2001
From: Jason Robinson <mail@jasonrobinson.me>
Date: Tue, 4 Jul 2017 23:17:19 +0300
Subject: [PATCH 3/7] Support receiving Diaspora new style encrypted JSON
 payloads

Decrypt the JSON and extract the Magic Envelope inside.

Closes #83
---
 CHANGELOG.md                               |  4 +++
 federation/protocols/diaspora/encrypted.py | 31 ++++++++++++++++++++
 federation/protocols/diaspora/protocol.py  | 33 +++++++++++++++++-----
 3 files changed, 61 insertions(+), 7 deletions(-)
 create mode 100644 federation/protocols/diaspora/encrypted.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 904a239..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))
diff --git a/federation/protocols/diaspora/encrypted.py b/federation/protocols/diaspora/encrypted.py
new file mode 100644
index 0000000..32fc9c5
--- /dev/null
+++ b/federation/protocols/diaspora/encrypted.py
@@ -0,0 +1,31 @@
+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 = json.loads(
+            cipher.decrypt(b64decode(payload.get("aes_key")), sentinel=None)
+        )
+        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/protocol.py b/federation/protocols/diaspora/protocol.py
index c0eb9df..11d9b11 100644
--- a/federation/protocols/diaspora/protocol.py
+++ b/federation/protocols/diaspora/protocol.py
@@ -2,6 +2,7 @@ import json
 import logging
 import warnings
 from base64 import b64decode, urlsafe_b64decode, b64encode, urlsafe_b64encode
+from json import JSONDecodeError
 from urllib.parse import unquote_plus
 
 from Crypto.Cipher import AES, PKCS1_v1_5
@@ -13,6 +14,7 @@ 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
 
 logger = logging.getLogger("federation")
@@ -55,17 +57,35 @@ 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(payload)
+            logger.debug("diaspora.protocol.store_magic_envelope_doc: json payload: %s", json_payload)
+            self.doc = self.get_json_payload_magic_envelope(json_payload)
+        except JSONDecodeError:
+            # 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)
+
+    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
@@ -83,7 +103,6 @@ class Protocol(BaseProtocol):
         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

From b19facc7af74735a91daf1fd56da94916d71484d Mon Sep 17 00:00:00 2001
From: Jason Robinson <mail@jasonrobinson.me>
Date: Tue, 4 Jul 2017 23:32:06 +0300
Subject: [PATCH 4/7] Remove an unused import

---
 federation/protocols/diaspora/protocol.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/federation/protocols/diaspora/protocol.py b/federation/protocols/diaspora/protocol.py
index 11d9b11..7f584ca 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 json import JSONDecodeError
 from urllib.parse import unquote_plus

From 92e7a957001e98dcf492adb3cc949aca31e749a6 Mon Sep 17 00:00:00 2001
From: Jason Robinson <mail@jasonrobinson.me>
Date: Tue, 4 Jul 2017 23:43:17 +0300
Subject: [PATCH 5/7] Catch ValueError instead of JSONDecodeError

Seems the latter was only added in Python 3.5 and we support 3.4 still.

Refs: #83
---
 federation/protocols/diaspora/protocol.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/federation/protocols/diaspora/protocol.py b/federation/protocols/diaspora/protocol.py
index 7f584ca..abc11e5 100644
--- a/federation/protocols/diaspora/protocol.py
+++ b/federation/protocols/diaspora/protocol.py
@@ -1,7 +1,6 @@
 import json
 import logging
 from base64 import b64decode, urlsafe_b64decode, b64encode, urlsafe_b64encode
-from json import JSONDecodeError
 from urllib.parse import unquote_plus
 
 from Crypto.Cipher import AES, PKCS1_v1_5
@@ -69,14 +68,15 @@ class Protocol(BaseProtocol):
         """Get the Magic Envelope, trying JSON first."""
         try:
             json_payload = json.loads(payload)
-            logger.debug("diaspora.protocol.store_magic_envelope_doc: json payload: %s", json_payload)
-            self.doc = self.get_json_payload_magic_envelope(json_payload)
-        except JSONDecodeError:
+        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.

From 670ec79fd545cfec5936eac3eef917e6debbd0e5 Mon Sep 17 00:00:00 2001
From: Jason Robinson <mail@jasonrobinson.me>
Date: Wed, 5 Jul 2017 00:56:28 +0300
Subject: [PATCH 6/7] Ensure JSON payloads are decoded from bytes when needed

Refs: #83
---
 federation/protocols/diaspora/encrypted.py | 5 ++---
 federation/protocols/diaspora/protocol.py  | 5 +++--
 federation/utils/text.py                   | 5 +++++
 3 files changed, 10 insertions(+), 5 deletions(-)
 create mode 100644 federation/utils/text.py

diff --git a/federation/protocols/diaspora/encrypted.py b/federation/protocols/diaspora/encrypted.py
index 32fc9c5..ccf23a4 100644
--- a/federation/protocols/diaspora/encrypted.py
+++ b/federation/protocols/diaspora/encrypted.py
@@ -20,9 +20,8 @@ class EncryptedPayload:
     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 = json.loads(
-            cipher.decrypt(b64decode(payload.get("aes_key")), sentinel=None)
-        )
+        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"))
diff --git a/federation/protocols/diaspora/protocol.py b/federation/protocols/diaspora/protocol.py
index abc11e5..73c417e 100644
--- a/federation/protocols/diaspora/protocol.py
+++ b/federation/protocols/diaspora/protocol.py
@@ -14,6 +14,7 @@ from federation.exceptions import EncryptedMessageError, NoSenderKeyFoundError,
 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:
@@ -67,7 +68,7 @@ class Protocol(BaseProtocol):
     def store_magic_envelope_doc(self, payload):
         """Get the Magic Envelope, trying JSON first."""
         try:
-            json_payload = json.loads(payload)
+            json_payload = json.loads(decode_if_bytes(payload))
         except ValueError:
             # XML payload
             xml = unquote_plus(payload)
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

From eed5e42f06595afc297f369e2258e60572094886 Mon Sep 17 00:00:00 2001
From: Jason Robinson <mail@jasonrobinson.me>
Date: Wed, 5 Jul 2017 16:36:59 +0300
Subject: [PATCH 7/7] Add tests for encrypted Diaspora JSON payload receive

Refs: #83
---
 .../protocols/diaspora/test_encrypted.py      | 33 +++++++++++++++++++
 .../tests/protocols/diaspora/test_protocol.py | 21 ++++++++++++
 federation/tests/utils/test_text.py           |  6 ++++
 3 files changed, 60 insertions(+)
 create mode 100644 federation/tests/protocols/diaspora/test_encrypted.py
 create mode 100644 federation/tests/utils/test_text.py

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"