package org.wildfly.swarm.jaxrs;

import java.util.Base64;

import javax.ws.rs.client.Client;
import javax.ws.rs.client.ClientBuilder;
import javax.ws.rs.client.Entity;
import javax.ws.rs.core.MediaType;
import javax.ws.rs.core.Response;

import org.jboss.arquillian.container.test.api.Deployment;
import org.jboss.arquillian.junit.Arquillian;
import org.jboss.shrinkwrap.api.Archive;
import org.jboss.shrinkwrap.api.ShrinkWrap;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.wildfly.swarm.jose.DecryptionOutput;
import org.wildfly.swarm.jose.Jose;
import org.wildfly.swarm.jose.JoseLookup;

@RunWith(Arquillian.class)
public class JoseCompactTest {

    private Client client;
    
    @Before
    public void setup() throws Exception {
        client = ClientBuilder.newClient();
    }
    
    @After
    public void tearDown() throws Exception {
        client.close();
    }
    
    @Deployment
    public static Archive<?> createDeployment() throws Exception {
        JAXRSArchive deployment = ShrinkWrap.create(JAXRSArchive.class);
        deployment.addResource(MyResource.class);
        deployment.addResource(SecuredApplication.class);
        deployment.addResource(JoseExceptionMapper.class); 
        deployment.addAllDependencies();
        deployment.addAsResource("keystore.jks", "keystore-compact.jks");
        deployment.addAsResource("project-jose-compact.yml", "project-defaults.yml");
        return deployment;
    }

    @Test
    public void testJwsCompact() throws Exception {
        Jose jose = JoseLookup.lookup().get();
        String signedData = client.target("http://localhost:8080/sign")
                                .request(MediaType.TEXT_PLAIN)
                                .post(Entity.entity(jose.sign("Hello"), MediaType.TEXT_PLAIN),
                                      String.class);
        Assert.assertEquals("Hello", jose.verify(signedData));
        Assert.assertEquals(3, signedData.split("\\.").length);
    }
    
    @Test
    public void testJwsCompactTampered() throws Exception {
        String[] jwsParts = JoseLookup.lookup().get().sign("Hello").split("\\.");
        Assert.assertEquals(3, jwsParts.length);
        String encodedContent = jwsParts[1];
        Assert.assertEquals("Hello", new String(Base64.getUrlDecoder().decode(encodedContent)));
        String newEncodedContent = new String(Base64.getUrlEncoder().encode("HellO".getBytes()));
        Assert.assertEquals("HellO", new String(Base64.getUrlDecoder().decode(newEncodedContent)));
        // Headers + content + signature
        String newJws = jwsParts[0] + "." + newEncodedContent + "." + jwsParts[2];  
        Assert.assertEquals(Response.Status.BAD_REQUEST.getStatusCode(), 
            client.target("http://localhost:8080/sign")
               .request(MediaType.TEXT_PLAIN)
               .post(Entity.entity(newJws, MediaType.TEXT_PLAIN))
               .getStatus());
    }
       
    @Test
    public void testJwsJweCompact() throws Exception {
        Jose jose = JoseLookup.lookup().get();
        String signedAndEncryptedData = client.target("http://localhost:8080/signAndEncrypt")
                                   .request(MediaType.TEXT_PLAIN)
                                   .post(Entity.entity(jose.encrypt(jose.sign("Hello")), MediaType.TEXT_PLAIN),
                                         String.class);
        String signedData = jose.decrypt(signedAndEncryptedData);
        Assert.assertEquals("Hello", jose.verify(signedData));
        
        Assert.assertEquals(3, signedData.split("\\.").length);
        Assert.assertEquals(5, signedAndEncryptedData.split("\\.").length);
    }
    
    @Test
    public void testJweCompact() throws Exception {
        Jose jose = JoseLookup.lookup().get();
        String encryptedData = client.target("http://localhost:8080/encrypt")
                                   .request(MediaType.TEXT_PLAIN)
                                   .post(Entity.entity(jose.encrypt("Hello"), MediaType.TEXT_PLAIN),
                                         String.class);
        DecryptionOutput decryption = jose.decryption(encryptedData);
        Assert.assertEquals("Hello", decryption.getData());
        Assert.assertEquals("server", decryption.getHeaders().get("kid"));
        Assert.assertEquals(5, encryptedData.split("\\.").length);
    }
    
    @Test
    public void testJweCompactTampered() throws Exception {
        Jose jose = JoseLookup.lookup().get();
        String[] jweParts = jose.encrypt("Hello").split("\\.");
        Assert.assertEquals(5, jweParts.length);
        String[] newJweParts = jose.encrypt("HellO").split("\\.");
        // Headers + IV + Encrypted CEK + Cipher + Authentication Tag 
        String newJwe = jweParts[0] + "." + jweParts[1] + "." + jweParts[2] + "."
                + newJweParts[3] + "." + jweParts[4]; 
        
        Assert.assertEquals(Response.Status.BAD_REQUEST.getStatusCode(), 
            client.target("http://localhost:8080/encrypt")
               .request(MediaType.TEXT_PLAIN)
               .post(Entity.entity(newJwe, MediaType.TEXT_PLAIN))
               .getStatus());
    }
}
