from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes

from server import app
import base64
import unittest
import os


class TestPosieService(unittest.TestCase):

    key_endpoint = "/key"
    decrypt_endpoint = "/decrypt"

    def setUp(self):
        self.key = os.urandom(32)
        self.iv = os.urandom(16)

        backend = default_backend()

        cipher = Cipher(algorithms.AES(self.key), modes.CTR(self.iv), backend=backend)

        self.encryptor = cipher.encryptor()

        # creates a test client
        self.app = app.test_client()

        # propagate the exceptions to the test client
        self.app.testing = True

        r = self.app.get(self.key_endpoint)

        key_string = base64.b64decode(r.data)

        self.public_key = serialization.load_der_public_key(
            key_string,
            backend=default_backend()
        )

    def send_message(self, message_bytes):
        data = self.encryptor.update(message_bytes) + self.encryptor.finalize()

        encrypted_key = self.public_key.encrypt(
            self.key,
            padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA1()),
                algorithm=hashes.SHA1(),
                label=None
            )
        )

        encoded_data = base64.b64encode(encrypted_key + self.iv + data)

        # Ask posie to decode message
        r = self.app.post(self.decrypt_endpoint, data=encoded_data)

        return r

    def test_key_generation(self):
        self.assertIsNotNone(self.public_key)

    def test_decrypt_fail_sends_400(self):

        # Ask posie to decode message
        r = self.app.post(self.decrypt_endpoint, data='rubbish')

        self.assertEqual(r.status_code, 400)

    def test_no_content_sends_400(self):

        # Ask posie to decode message
        r = self.app.post(self.decrypt_endpoint, data='')

        self.assertEqual(r.status_code, 400)

    def test_decrypts_message(self):
        # Encrypt a message with the key
        message = b"Some encrypted message"

        # Ask posie to decode message
        r = self.send_message(message)

        # Compare to bytestring version of decrypted data
        self.assertEqual(r.data, message)

    def test_decrypts_large_message(self):
        # Encrypt a message with the key
        message = b"Lorem ipsum dolor sit amet, consectetur adipiscing elit. Ut lectus urna, elementum vitae massa congue, varius aliquam purus. Etiam quis pellentesque felis. Morbi vehicula sagittis enim iaculis pretium. Fusce venenatis pretium ante, et semper ipsum blandit vitae. Lorem ipsum dolor sit amet, consectetur adipiscing elit. Cras lobortis dictum sapien et faucibus. Nulla vitae nisl ante."

        # Ask posie to decode message
        r = self.send_message(message)

        # Compare to bytestring version of decrypted data
        self.assertEqual(r.data, message)

@@ -2508,393 +2508,1008 @@
  = b
-%22Lorem ipsum dolor sit amet, consectetur adipiscing elit. Ut lectus urna, elementum vitae massa congue, varius aliquam purus. Etiam quis pellentesque felis. Morbi vehicula sagittis enim iaculis pretium. Fusce venenatis pretium ante, et semper ipsum blandit vitae. Lorem ipsum dolor sit amet, consectetur adipiscing elit. Cras lobortis dictum sapien et faucibus. Nulla vitae nisl ante.%22
+'''%7B%0A            %22type%22: %22uk.gov.ons.edc.eq:surveyresponse%22,%0A            %22version%22: %220.0.1%22,%0A            %22origin%22: %22uk.gov.ons.edc.eq%22,%0A            %22survey_id%22: %2221%22,%0A            %22collection%22: %7B%0A              %22exercise_sid%22: %22hfjdskf%22,%0A              %22instrument_id%22: %220203%22,%0A              %22period%22: %222016-02-01%22%0A            %7D,%0A            %22submitted_at%22: %222016-03-12T10:39:40Z%22,%0A            %22metadata%22: %7B%0A              %22user_id%22: %22789473423%22,%0A              %22ru_ref%22: %2212345678901A%22%0A            %7D,%0A            %22data%22: %7B%0A              %2211%22: %2201042016%22,%0A              %2212%22: %2231102016%22,%0A              %2220%22: %221800000%22,%0A              %2251%22: %2284%22,%0A              %2252%22: %2210%22,%0A              %2253%22: %2273%22,%0A              %2254%22: %2224%22,%0A              %2250%22: %22205%22,%0A              %2222%22: %22705000%22,%0A              %2223%22: %22900%22,%0A              %2224%22: %2274%22,%0A              %2225%22: %2250%22,%0A              %2226%22: %22100%22,%0A              %2221%22: %2260000%22,%0A              %2227%22: %227400%22,%0A              %22146%22: %22some comment%22%0A            %7D%0A        %7D'''
 %0A%0A  

Use some json for large message test