Repository: jaywink/federation

Issue:
Add support for Diaspora Salmon refactoring
See diaspora/diaspora_federation#30

Discussion:
pep8speaks: Hello @jaywink! Thanks for submitting the PR.



 - In the file [`federation/protocols/diaspora/magic_envelope.py`](https://github.com/jaywink/federation/blob/54dc84293126a9c0330ce7b0efb08bada5ed22bf/federation/protocols/diaspora/magic_envelope.py), following are the PEP8 issues :

> [Line 37:1](https://github.com/jaywink/federation/blob/54dc84293126a9c0330ce7b0efb08bada5ed22bf/federation/protocols/diaspora/magic_envelope.py#L37): [W293](https://duckduckgo.com/?q=pep8%20W293) blank line contains whitespace





 - In the file [`federation/tests/protocols/diaspora/test_protocol.py`](https://github.com/jaywink/federation/blob/54dc84293126a9c0330ce7b0efb08bada5ed22bf/federation/tests/protocols/diaspora/test_protocol.py), following are the PEP8 issues :

> [Line 145:45](https://github.com/jaywink/federation/blob/54dc84293126a9c0330ce7b0efb08bada5ed22bf/federation/tests/protocols/diaspora/test_protocol.py#L145): [E711](https://duckduckgo.com/?q=pep8%20E711) comparison to None should be 'if cond is None:'



Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 64f576a..2d5de32 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,13 @@
 # Changelog
 
+## [unreleased]
+
+### Backwards incompatible changes
+* Removed exception class `NoHeaderInMessageError`. New style Diaspora protocol does not have a custom header in the Salmon magic envelope and thus there is no need to raise this anywhere.
+
+### Added
+* New style Diaspora public payloads are now supported (see [here](https://github.com/diaspora/diaspora_federation/issues/30)). Old style payloads are still supported. Payloads are also still sent out old style.
+
 ## [0.11.0] - 2017-05-08
 
 ### Backwards incompatible changes
diff --git a/federation/exceptions.py b/federation/exceptions.py
index f94dd49..243111b 100644
--- a/federation/exceptions.py
+++ b/federation/exceptions.py
@@ -3,11 +3,6 @@ class EncryptedMessageError(Exception):
     pass
 
 
-class NoHeaderInMessageError(Exception):
-    """Message payload is missing required header."""
-    pass
-
-
 class NoSenderKeyFoundError(Exception):
     """Sender private key was not available to sign a payload message."""
     pass
diff --git a/federation/protocols/diaspora/magic_envelope.py b/federation/protocols/diaspora/magic_envelope.py
index e3f6fff..ff9cc13 100644
--- a/federation/protocols/diaspora/magic_envelope.py
+++ b/federation/protocols/diaspora/magic_envelope.py
@@ -1,18 +1,21 @@
-from base64 import urlsafe_b64encode, b64encode
+from base64 import urlsafe_b64encode, b64encode, urlsafe_b64decode
 
 from Crypto.Hash import SHA256
 from Crypto.Signature import PKCS1_v1_5 as PKCSSign
 from lxml import etree
 
 
-class MagicEnvelope(object):
+NAMESPACE = "http://salmon-protocol.org/ns/magic-env"
+
+
+class MagicEnvelope():
     """Diaspora protocol magic envelope.
 
     See: http://diaspora.github.io/diaspora_federation/federation/magicsig.html
     """
 
     nsmap = {
-        'me': 'http://salmon-protocol.org/ns/magic-env'
+        "me": NAMESPACE,
     }
 
     def __init__(self, message, private_key, author_handle, wrap_payload=False):
@@ -28,6 +31,16 @@ class MagicEnvelope(object):
         self.doc = None
         self.payload = None
 
+    @staticmethod
+    def get_sender(doc):
+        """Get the key_id from the `sig` element which contains urlsafe_b64encoded Diaspora handle.
+
+        :param doc: ElementTree document
+        :returns: Diaspora handle
+        """
+        key_id = doc.find(".//{%s}sig" % NAMESPACE).get("key_id")
+        return urlsafe_b64decode(key_id).decode("utf-8")
+
     def create_payload(self):
         """Create the payload doc.
 
@@ -58,14 +71,13 @@ class MagicEnvelope(object):
         return sig, key_id
 
     def build(self):
-        self.doc = etree.Element("{%s}env" % self.nsmap["me"], nsmap=self.nsmap)
-        etree.SubElement(self.doc, "{%s}encoding" % self.nsmap["me"]).text = 'base64url'
-        etree.SubElement(self.doc, "{%s}alg" % self.nsmap["me"]).text = 'RSA-SHA256'
+        self.doc = etree.Element("{%s}env" % NAMESPACE, nsmap=self.nsmap)
+        etree.SubElement(self.doc, "{%s}encoding" % NAMESPACE).text = 'base64url'
+        etree.SubElement(self.doc, "{%s}alg" % NAMESPACE).text = 'RSA-SHA256'
         self.create_payload()
-        etree.SubElement(self.doc, "{%s}data" % self.nsmap["me"],
-                         {"type": "application/xml"}).text = self.payload
+        etree.SubElement(self.doc, "{%s}data" % NAMESPACE, {"type": "application/xml"}).text = self.payload
         signature, key_id = self._build_signature()
-        etree.SubElement(self.doc, "{%s}sig" % self.nsmap["me"], key_id=key_id).text = signature
+        etree.SubElement(self.doc, "{%s}sig" % NAMESPACE, key_id=key_id).text = signature
         return self.doc
 
     def render(self):
diff --git a/federation/protocols/diaspora/protocol.py b/federation/protocols/diaspora/protocol.py
index 8113b90..c6bce54 100644
--- a/federation/protocols/diaspora/protocol.py
+++ b/federation/protocols/diaspora/protocol.py
@@ -11,10 +11,9 @@ from Crypto.Random import get_random_bytes
 from Crypto.Signature import PKCS1_v1_5 as PKCSSign
 from lxml import etree
 
-from federation.exceptions import (
-    EncryptedMessageError, NoHeaderInMessageError, NoSenderKeyFoundError, SignatureVerificationError,
-)
+from federation.exceptions import EncryptedMessageError, NoSenderKeyFoundError, SignatureVerificationError
 from federation.protocols.base import BaseProtocol
+from federation.protocols.diaspora.magic_envelope import MagicEnvelope
 
 logger = logging.getLogger("federation")
 
@@ -67,6 +66,7 @@ class Protocol(BaseProtocol):
         xml = xml.lstrip().encode("utf-8")
         logger.debug("diaspora.protocol.receive: xml content: %s", xml)
         self.doc = etree.fromstring(xml)
+        # Check for a legacy header
         self.find_header()
         # Open payload and get actual message
         self.content = self.get_message_content()
@@ -88,12 +88,16 @@ 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:
-            self.encrypted = False
+            # Legacy public header found
+            self.legacy = True
             return
         if self.doc.find(".//{" + PROTOCOL_NS + "}encrypted_header") == None:
-            raise NoHeaderInMessageError("Could not find header in message")
+            # No legacy encrypted header found
+            return
+        self.legacy = True
         if not self.user:
             raise EncryptedMessageError("Cannot decrypt private message without user object")
         user_private_key = self._get_user_key(self.user)
@@ -104,6 +108,11 @@ class Protocol(BaseProtocol):
         )
 
     def get_sender(self):
+        if self.legacy:
+            return self.get_sender_legacy()
+        return MagicEnvelope.get_sender(self.doc)
+
+    def get_sender_legacy(self):
         try:
             return self.header.find(".//{"+PROTOCOL_NS+"}author_id").text
         except AttributeError:


Test patch:
diff --git a/federation/tests/protocols/diaspora/test_magic_envelope.py b/federation/tests/protocols/diaspora/test_magic_envelope.py
index f9bd925..952ec2e 100644
--- a/federation/tests/protocols/diaspora/test_magic_envelope.py
+++ b/federation/tests/protocols/diaspora/test_magic_envelope.py
@@ -1,12 +1,15 @@
+from lxml import etree
+
 from Crypto import Random
 from Crypto.PublicKey import RSA
 from lxml.etree import _Element
 
 from federation.protocols.diaspora.magic_envelope import MagicEnvelope
 from federation.tests.fixtures.keys import get_dummy_private_key
+from federation.tests.fixtures.payloads import DIASPORA_PUBLIC_PAYLOAD
 
 
-class TestMagicEnvelope(object):
+class TestMagicEnvelope():
     @staticmethod
     def generate_rsa_private_key():
         """Generate a new RSA private key."""
@@ -80,3 +83,7 @@ class TestMagicEnvelope(object):
         )
         output2 = env2.render()
         assert output2 == output
+
+    def test_get_sender(self):
+        doc = etree.fromstring(bytes(DIASPORA_PUBLIC_PAYLOAD, encoding="utf-8"))
+        assert MagicEnvelope.get_sender(doc) == "foobar@example.com"
diff --git a/federation/tests/protocols/diaspora/test_protocol.py b/federation/tests/protocols/diaspora/test_protocol.py
index fa4359e..3fd5a37 100644
--- a/federation/tests/protocols/diaspora/test_protocol.py
+++ b/federation/tests/protocols/diaspora/test_protocol.py
@@ -5,7 +5,7 @@ from xml.etree.ElementTree import ElementTree
 from lxml import etree
 import pytest
 
-from federation.exceptions import EncryptedMessageError, NoSenderKeyFoundError, NoHeaderInMessageError
+from federation.exceptions import EncryptedMessageError, NoSenderKeyFoundError
 from federation.protocols.diaspora.protocol import Protocol, identify_payload
 from federation.tests.fixtures.payloads import (
     ENCRYPTED_LEGACY_DIASPORA_PAYLOAD, UNENCRYPTED_LEGACY_DIASPORA_PAYLOAD,
@@ -105,12 +105,6 @@ class TestDiasporaProtocol(DiasporaTestBase):
         with pytest.raises(NoSenderKeyFoundError):
             protocol.receive(UNENCRYPTED_LEGACY_DIASPORA_PAYLOAD, user, mock_not_found_get_contact_key)
 
-    def test_find_header_raises_if_header_cannot_be_found(self):
-        protocol = self.init_protocol()
-        protocol.doc = etree.fromstring("<foo>bar</foo>")
-        with pytest.raises(NoHeaderInMessageError):
-            protocol.find_header()
-
     def test_get_message_content(self):
         protocol = self.init_protocol()
         protocol.doc = self.get_unencrypted_doc()
@@ -130,25 +124,25 @@ class TestDiasporaProtocol(DiasporaTestBase):
     def test_identify_payload_with_other_payload(self):
         assert identify_payload("foobar not a diaspora protocol") == False
 
-    def test_get_sender_returns_sender_in_header(self):
+    def test_get_sender_legacy_returns_sender_in_header(self):
         protocol = self.init_protocol()
         protocol.doc = self.get_unencrypted_doc()
         protocol.find_header()
-        assert protocol.get_sender() == "bob@example.com"
+        assert protocol.get_sender_legacy() == "bob@example.com"
 
-    def test_get_sender_returns_sender_in_content(self):
+    def test_get_sender_legacy_returns_sender_in_content(self):
         protocol = self.init_protocol()
         protocol.header = ElementTree()
         protocol.content = "<content><diaspora_handle>bob@example.com</diaspora_handle></content>"
-        assert protocol.get_sender() == "bob@example.com"
+        assert protocol.get_sender_legacy() == "bob@example.com"
         protocol.content = "<content><sender_handle>bob@example.com</sender_handle></content>"
-        assert protocol.get_sender() == "bob@example.com"
+        assert protocol.get_sender_legacy() == "bob@example.com"
 
-    def test_get_sender_returns_none_if_no_sender_found(self):
+    def test_get_sender_legacy_returns_none_if_no_sender_found(self):
         protocol = self.init_protocol()
         protocol.header = ElementTree()
         protocol.content = "<content><handle>bob@example.com</handle></content>"
-        assert protocol.get_sender() == None
+        assert protocol.get_sender_legacy() is None
 
     @patch.object(Protocol, "init_message")
     @patch.object(Protocol, "create_salmon_envelope")