Add more encryption tests

tests/test_encryption.py

@@ -1,6 +1,22 @@
+import os
 import unittest
 import hashlib
-from minecraft.networking.encryption import minecraft_sha1_hash_digest
+from io import BytesIO
+from minecraft.networking.encryption import (
+    minecraft_sha1_hash_digest,
+    encrypt_token_and_secret,
+    generate_shared_secret,
+    generate_verification_hash,
+    create_AES_cipher,
+    EncryptedFileObjectWrapper,
+)
+
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives.asymmetric.padding import PKCS1v15
+from cryptography.hazmat.primitives.serialization import load_der_private_key
+
+KEY_LOCATION = os.path.join(os.path.dirname(os.path.realpath(__file__)),
+                            "encryption")
 
 
 class Hashing(unittest.TestCase):
@@ -13,3 +29,51 @@ class Hashing(unittest.TestCase):
             sha1_hash = hashlib.sha1()
             sha1_hash.update(input_value.encode('utf-8'))
             self.assertEquals(minecraft_sha1_hash_digest(sha1_hash), result)
+
+
+class Encryption(unittest.TestCase):
+
+    def setUp(self):
+        with open(os.path.join(KEY_LOCATION, "priv_key.bin"), "rb") as f:
+            self.private_key = f.read()
+        self.private_key = load_der_private_key(self.private_key, None,
+                                                default_backend())
+
+        with open(os.path.join(KEY_LOCATION, "pub_key.bin"), "rb") as f:
+            self.public_key = f.read()
+        self.token = generate_shared_secret()
+
+    def test_token_secret_encryption(self):
+        secret = generate_shared_secret()
+        token, encrypted_secret = encrypt_token_and_secret(self.public_key,
+                                                           self.token, secret)
+        decrypted_token = self.private_key.decrypt(token,
+                                                   PKCS1v15())
+        decrypted_secret = self.private_key.decrypt(encrypted_secret,
+                                                    PKCS1v15())
+
+        self.assertEquals(self.token, decrypted_token)
+        self.assertEquals(secret, decrypted_secret)
+
+    def generate_hash_test(self):
+        verification_hash = generate_verification_hash(
+            b"", "secret".encode('utf-8'),  self.public_key)
+        self.assertEquals("1f142e737a84a974a5f2a22f6174a78d80fd97f5",
+                          verification_hash)
+
+    def test_file_object_wrapper(self):
+        cipher = create_AES_cipher(generate_shared_secret())
+        encryptor = cipher.encryptor()
+        decryptor = cipher.decryptor()
+
+        test_data = "hello".encode('utf-8')
+        io = BytesIO()
+        io.write(encryptor.update(test_data))
+        io.seek(0)
+
+        file_object_wrapper = EncryptedFileObjectWrapper(io, decryptor)
+        decrypted_data = file_object_wrapper.read(len(test_data))
+
+        self.assertEqual(test_data, decrypted_data)
+
+    # TODO: test for the socket wrapper