package com.jenjinstudios.core.message;

import com.jenjinstudios.core.Connection;
import com.jenjinstudios.core.MessageIO;
import com.jenjinstudios.core.io.Message;
import com.jenjinstudios.core.io.MessageOutputStream;
import org.testng.annotations.Test;

import java.net.InetAddress;
import java.security.Key;
import java.security.KeyPair;
import java.util.HashMap;
import java.util.Map;

import static org.mockito.Mockito.*;

/**
 * Test key verification.
 *
 * @author Caleb Brinkman
 */
public class ExecutablePublicKeyMessageTest
{
    /**
     * Test key verification.
     */
    @Test
    public void testVerification() {
        KeyPair rsaKeyPair = Connection.generateRSAKeyPair();
        Message message = Connection.generatePublicKeyMessage(rsaKeyPair.getPublic());
        InetAddress address = InetAddress.getLoopbackAddress();
        Map<InetAddress, Key> keys = new HashMap<>(10);
        keys.put(address, rsaKeyPair.getPublic());

        Connection connection = mock(Connection.class);
        MessageIO messageIO = mock(MessageIO.class);
        MessageOutputStream out = mock(MessageOutputStream.class);

        when(messageIO.getOut()).thenReturn(out);
        when(messageIO.getAddress()).thenReturn(address);
        when(connection.getMessageIO()).thenReturn(messageIO);
        when(connection.getVerifiedKeys()).thenReturn(keys);

        ExecutablePublicKeyMessage executable = new ExecutablePublicKeyMessage(connection, message);
        executable.runImmediate();

        verify(out).setPublicKey(rsaKeyPair.getPublic());
    }

    /**
     * Test verification bypass when no verified keys are present.
     */
    @Test
    public void testVerificationNoKeys() {
        KeyPair rsaKeyPair = Connection.generateRSAKeyPair();
        Message message = Connection.generatePublicKeyMessage(rsaKeyPair.getPublic());
        InetAddress address = InetAddress.getLoopbackAddress();
        Map<InetAddress, Key> keys = new HashMap<>(10);

        Connection connection = mock(Connection.class);
        MessageIO messageIO = mock(MessageIO.class);
        MessageOutputStream out = mock(MessageOutputStream.class);

        when(messageIO.getOut()).thenReturn(out);
        when(messageIO.getAddress()).thenReturn(address);
        when(connection.getMessageIO()).thenReturn(messageIO);
        when(connection.getVerifiedKeys()).thenReturn(keys);

        ExecutablePublicKeyMessage executable = new ExecutablePublicKeyMessage(connection, message);
        executable.runImmediate();

        verify(out).setPublicKey(rsaKeyPair.getPublic());
    }

    /**
     * Test failed key verification due to invalid key.
     */
    @Test
    public void testVerificationInvalidKey() {
        KeyPair rsaKeyPair = Connection.generateRSAKeyPair();
        KeyPair invalidKeyPair = Connection.generateRSAKeyPair();
        Message message = Connection.generatePublicKeyMessage(invalidKeyPair.getPublic());
        InetAddress address = InetAddress.getLoopbackAddress();
        Map<InetAddress, Key> keys = new HashMap<>(10);
        keys.put(address, rsaKeyPair.getPublic());

        Connection connection = mock(Connection.class);
        MessageIO messageIO = mock(MessageIO.class);
        MessageOutputStream out = mock(MessageOutputStream.class);

        when(messageIO.getOut()).thenReturn(out);
        when(messageIO.getAddress()).thenReturn(address);
        when(connection.getMessageIO()).thenReturn(messageIO);
        when(connection.getVerifiedKeys()).thenReturn(keys);

        ExecutablePublicKeyMessage executable = new ExecutablePublicKeyMessage(connection, message);
        executable.runImmediate();

        verify(out, times(0)).setPublicKey(any());
    }

    /**
     * Test invalid key verification due to unknown host.
     */
    @Test
    public void testVerificationNoAddress() {
        KeyPair rsaKeyPair = Connection.generateRSAKeyPair();
        KeyPair invalidKeyPair = Connection.generateRSAKeyPair();
        Message message = Connection.generatePublicKeyMessage(invalidKeyPair.getPublic());
        InetAddress address = InetAddress.getLoopbackAddress();
        Map<InetAddress, Key> keys = new HashMap<>(10);
        keys.put(address, rsaKeyPair.getPublic());

        Connection connection = mock(Connection.class);
        MessageIO messageIO = mock(MessageIO.class);
        MessageOutputStream out = mock(MessageOutputStream.class);

        when(messageIO.getOut()).thenReturn(out);
        when(connection.getMessageIO()).thenReturn(messageIO);
        when(connection.getVerifiedKeys()).thenReturn(keys);

        ExecutablePublicKeyMessage executable = new ExecutablePublicKeyMessage(connection, message);
        executable.runImmediate();

        verify(out, times(0)).setPublicKey(any());
    }
}

@@ -159,16 +159,69 @@
 essage;%0A
+import com.jenjinstudios.core.io.MessageInputStream;%0A
 import c
@@ -1066,35 +1066,37 @@
     MessageI
-O messageIO
+nputStream in
  = mock(Mess
@@ -1091,33 +1091,42 @@
  = mock(MessageI
-O
+nputStream
 .class);%0A       
@@ -1190,38 +1190,33 @@
 );%0A%0A        
-when(m
+M
 essageIO
 .getOut()).t
@@ -1207,84 +1207,44 @@
 geIO
-.getOut()).thenReturn(out);%0A        when(messageIO.getAddress()).thenReturn(
+ messageIO = new MessageIO(in, out, 
 addr
@@ -1241,32 +1241,33 @@
  out, address);%0A
+%0A
         when(con

Test tweakage, brah