package org.matrix.androidsdk;

import android.os.MemoryFile;
import android.support.test.runner.AndroidJUnit4;
import android.text.TextUtils;
import android.util.Base64;

import org.junit.Assert;
import org.junit.FixMethodOrder;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.MethodSorters;
import org.matrix.androidsdk.crypto.MXEncryptedAttachments;
import org.matrix.androidsdk.rest.model.crypto.EncryptedFileInfo;
import org.matrix.androidsdk.rest.model.crypto.EncryptedFileKey;

import java.io.InputStream;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;

/**
 * Unit tests AttachmentEncryptionTest.
 */
@RunWith(AndroidJUnit4.class)
@FixMethodOrder(MethodSorters.NAME_ASCENDING)
public class AttachmentEncryptionTest {

    private String checkDecryption(String input, EncryptedFileInfo encryptedFileInfo) throws Exception  {
        byte[] in = Base64.decode(input, Base64.DEFAULT);

        MemoryFile memoryFile;

        if (0 == in.length) {
            memoryFile = new MemoryFile("file" + System.currentTimeMillis(), 0);
        } else {
            memoryFile = new MemoryFile("file" + System.currentTimeMillis(), in.length);
            memoryFile.getOutputStream().write(in);
        }

        InputStream decryptedStream = MXEncryptedAttachments.decryptAttachment(memoryFile.getInputStream(), encryptedFileInfo);

        Assert.assertTrue(null != decryptedStream);

        byte[] buffer = new byte[100];

        int len = decryptedStream.read(buffer);

        return Base64.encodeToString(buffer, 0, len, Base64.DEFAULT).replaceAll("\n", "").replaceAll("=", "");
    }

    @Test
    public void checkDecrypt1() throws Exception {
        EncryptedFileInfo encryptedFileInfo = new EncryptedFileInfo();
        encryptedFileInfo.v = "v1";
        Map<String, String> hashMap = new HashMap<>();
        hashMap.put("sha256", "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU");
        encryptedFileInfo.hashes = hashMap;

        encryptedFileInfo.key = new EncryptedFileKey();
        encryptedFileInfo.key.alg =  "A256CTR";
        encryptedFileInfo.key.k = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
        encryptedFileInfo.key.key_ops = Arrays.asList("encrypt","decrypt");
        encryptedFileInfo.key.kty = "oct";

        encryptedFileInfo.iv = "AAAAAAAAAAAAAAAAAAAAAA";

        Assert.assertTrue(TextUtils.equals(checkDecryption("", encryptedFileInfo), ""));
    }

    @Test
    public void checkDecrypt2() throws Exception {
        EncryptedFileInfo encryptedFileInfo = new EncryptedFileInfo();
        encryptedFileInfo.v = "v1";
        Map<String, String> hashMap = new HashMap<>();
        hashMap.put("sha256", "YzF08lARDdOCzJpzuSwsjTNlQc4pHxpdHcXiD/wpK6k");
        encryptedFileInfo.hashes = hashMap;

        encryptedFileInfo.key = new EncryptedFileKey();
        encryptedFileInfo.key.alg =  "A256CTR";
        encryptedFileInfo.key.k = "__________________________________________8";
        encryptedFileInfo.key.key_ops = Arrays.asList("encrypt","decrypt");
        encryptedFileInfo.key.kty = "oct";

        encryptedFileInfo.iv = "//////////8AAAAAAAAAAA";

        Assert.assertTrue(TextUtils.equals(checkDecryption("5xJZTt5cQicm+9f4", encryptedFileInfo), "SGVsbG8sIFdvcmxk"));
    }

    @Test
    public void checkDecrypt3() throws Exception {
        EncryptedFileInfo encryptedFileInfo = new EncryptedFileInfo();
        encryptedFileInfo.v = "v2";
        Map<String, String> hashMap = new HashMap<>();
        hashMap.put("sha256", "IOq7/dHHB+mfHfxlRY5XMeCWEwTPmlf4cJcgrkf6fVU");
        encryptedFileInfo.hashes = hashMap;

        encryptedFileInfo.key = new EncryptedFileKey();
        encryptedFileInfo.key.alg = "A256CTR";
        encryptedFileInfo.key.k = "__________________________________________8";
        encryptedFileInfo.key.key_ops = Arrays.asList("encrypt","decrypt");
        encryptedFileInfo.key.kty = "oct";

        encryptedFileInfo.iv = "//////////8AAAAAAAAAAA";

        Assert.assertTrue(TextUtils.equals(checkDecryption(
                "zhtFStAeFx0s+9L/sSQO+WQMtldqYEHqTxMduJrCIpnkyer09kxJJuA4K+adQE4w+7jZe/vR9kIcqj9rOhDR8Q",
                encryptedFileInfo),
                "YWxwaGFudW1lcmljYWxseWFscGhhbnVtZXJpY2FsbHlhbHBoYW51bWVyaWNhbGx5YWxwaGFudW1lcmljYWxseQ"));
    }

    @Test
    public void checkDecrypt4() throws Exception {
        EncryptedFileInfo encryptedFileInfo = new EncryptedFileInfo();
        encryptedFileInfo.v = "v1";
        Map<String, String> hashMap = new HashMap<>();
        hashMap.put("sha256", "LYG/orOViuFwovJpv2YMLSsmVKwLt7pY3f8SYM7KU5E");
        encryptedFileInfo.hashes = hashMap;

        encryptedFileInfo.key = new EncryptedFileKey();
        encryptedFileInfo.key.alg =  "A256CTR";
        encryptedFileInfo.key.k = "__________________________________________8";
        encryptedFileInfo.key.key_ops = Arrays.asList("encrypt","decrypt");
        encryptedFileInfo.key.kty = "oct";

        encryptedFileInfo.iv = "/////////////////////w";

        Assert.assertTrue(!TextUtils.equals(checkDecryption(
                "tJVNBVJ/vl36UQt4Y5e5m84bRUrQHhcdLPvS/7EkDvlkDLZXamBB6k8THbiawiKZ5Mnq9PZMSSbgOCvmnUBOMA",
                encryptedFileInfo),
                "YWxwaGFudW1lcmljYWxseWFscGhhbnVtZXJpY2FsbHlhbHBoYW51bWVyaWNhbGx5YWxwaGFudW1lcmljYWxseQ"));
    }
}

@@ -511,16 +511,53 @@
 leKey;%0A%0A
+import java.io.ByteArrayInputStream;%0A
 import j
@@ -992,29 +992,31 @@
     
-MemoryFile memoryFile
+InputStream inputStream
 ;%0A%0A 
@@ -1060,74 +1060,49 @@
     
-memoryFile = new MemoryFile(%22file%22 + System.currentTimeMillis(), 0
+inputStream = new ByteArrayInputStream(in
 );%0A 
@@ -1121,32 +1121,43 @@
 e %7B%0A            
+MemoryFile 
 memoryFile = new
@@ -1265,24 +1265,79 @@
 .write(in);%0A
+            inputStream = memoryFile.getInputStream();%0A
         %7D%0A%0A 
@@ -1414,31 +1414,17 @@
 achment(
-memoryFile.getI
+i
 nputStre
@@ -1425,18 +1425,16 @@
 utStream
-()
 , encryp

Fix checkDecrypt1(): It is not possible to create a 0-length MemoryFile. Use a ByteArrayInputStream instead