# Path: federation/protocols/diaspora/encrypted.py
# def pkcs7_unpad(data):
#     """
#     Remove the padding bytes that were added at point of encryption.
# 
#     Implementation copied from pyaspora:
#     https://github.com/mjnovice/pyaspora/blob/master/pyaspora/diaspora/protocol.py#L209
#     """
#     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))
# 
#     @staticmethod
#     def get_aes_key_json(iv, key):
#         return json.dumps({
#             "key": b64encode(key).decode("ascii"),
#             "iv": b64encode(iv).decode("ascii"),
#         }).encode("utf-8")
# 
#     @staticmethod
#     def get_iv_key_encrypter():
#         iv = get_random_bytes(AES.block_size)
#         key = get_random_bytes(32)
#         encrypter = AES.new(key, AES.MODE_CBC, iv)
#         return iv, key, encrypter
# 
#     @staticmethod
#     def encrypt(payload, public_key):
#         """
#         Encrypt a payload using an encrypted JSON wrapper.
# 
#         See: https://diaspora.github.io/diaspora_federation/federation/encryption.html
# 
#         :param payload: Payload document as a string.
#         :param public_key: Public key of recipient as an RSA object.
#         :return: Encrypted JSON wrapper as dict.
#         """
#         iv, key, encrypter = EncryptedPayload.get_iv_key_encrypter()
#         aes_key_json = EncryptedPayload.get_aes_key_json(iv, key)
#         cipher = PKCS1_v1_5.new(public_key)
#         aes_key = b64encode(cipher.encrypt(aes_key_json))
#         padded_payload = pkcs7_pad(payload.encode("utf-8"), AES.block_size)
#         encrypted_me = b64encode(encrypter.encrypt(padded_payload))
#         return {
#             "aes_key": aes_key.decode("utf-8"),
#             "encrypted_magic_envelope": encrypted_me.decode("utf8"),
#         }
# 
# Path: federation/tests/fixtures/keys.py
# def get_dummy_private_key():
#     return RSA.importKey(PRIVATE_KEY)


from unittest.mock import patch, Mock
from Crypto.Cipher import AES
from lxml import etree
from federation.protocols.diaspora.encrypted import pkcs7_unpad, EncryptedPayload
from federation.tests.fixtures.keys import get_dummy_private_key

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"

    def test_encrypt(self):
        private_key = get_dummy_private_key()
        public_key = private_key.publickey()
        encrypted = EncryptedPayload.encrypt("<spam>eggs</spam>", public_key)
        assert "aes_key" in encrypted
        assert "encrypted_magic_envelope" in encrypted
        # See we can decrypt it too

# Path: federation/protocols/diaspora/encrypted.py
# def pkcs7_unpad(data):
#     """
#     Remove the padding bytes that were added at point of encryption.
# 
#     Implementation copied from pyaspora:
#     https://github.com/mjnovice/pyaspora/blob/master/pyaspora/diaspora/protocol.py#L209
#     """
#     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))
# 
#     @staticmethod
#     def get_aes_key_json(iv, key):
#         return json.dumps({
#             "key": b64encode(key).decode("ascii"),
#             "iv": b64encode(iv).decode("ascii"),
#         }).encode("utf-8")
# 
#     @staticmethod
#     def get_iv_key_encrypter():
#         iv = get_random_bytes(AES.block_size)
#         key = get_random_bytes(32)
#         encrypter = AES.new(key, AES.MODE_CBC, iv)
#         return iv, key, encrypter
# 
#     @staticmethod
#     def encrypt(payload, public_key):
#         """
#         Encrypt a payload using an encrypted JSON wrapper.
# 
#         See: https://diaspora.github.io/diaspora_federation/federation/encryption.html
# 
#         :param payload: Payload document as a string.
#         :param public_key: Public key of recipient as an RSA object.
#         :return: Encrypted JSON wrapper as dict.
#         """
#         iv, key, encrypter = EncryptedPayload.get_iv_key_encrypter()
#         aes_key_json = EncryptedPayload.get_aes_key_json(iv, key)
#         cipher = PKCS1_v1_5.new(public_key)
#         aes_key = b64encode(cipher.encrypt(aes_key_json))
#         padded_payload = pkcs7_pad(payload.encode("utf-8"), AES.block_size)
#         encrypted_me = b64encode(encrypter.encrypt(padded_payload))
#         return {
#             "aes_key": aes_key.decode("utf-8"),
#             "encrypted_magic_envelope": encrypted_me.decode("utf8"),
#         }
# 
# Path: federation/tests/fixtures/keys.py
# def get_dummy_private_key():
#     return RSA.importKey(PRIVATE_KEY)


# Path: federation/tests/protocols/diaspora/test_encrypted.py
from unittest.mock import patch, Mock
from Crypto.Cipher import AES
from lxml import etree
from federation.protocols.diaspora.encrypted import pkcs7_unpad, EncryptedPayload
from federation.tests.fixtures.keys import get_dummy_private_key


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"

    def test_encrypt(self):
        private_key = get_dummy_private_key()
        public_key = private_key.publickey()
        encrypted = EncryptedPayload.encrypt("<spam>eggs</spam>", public_key)
        assert "aes_key" in encrypted
        assert "encrypted_magic_envelope" in encrypted
        # See we can decrypt it too

        decrypted = EncryptedPayload.decrypt(encrypted, private_key)