package net.lshift.spki.suiteb.demo;

import static org.hamcrest.CoreMatchers.is;
import static org.junit.Assert.assertThat;

import java.io.IOException;
import java.io.PrintWriter;

import net.lshift.spki.InvalidInputException;
import net.lshift.spki.ParseException;
import net.lshift.spki.PrettyPrinter;
import net.lshift.spki.convert.openable.ByteOpenable;
import org.junit.Test;

/**
 * You can consider the first part of each test the setting up of the system,
 * and key management. Then, the call to sendMessageFromServerToClient is the
 * system in actual operation.
 */
public class TestLoop {
    @Test
    public void serverSendsPlainTextMessageToClientThatTrustsIt()
            throws IOException, InvalidInputException {
        Server server = new Server();
        Client client = new Client(server.getPublicSigningKey().getKeyId());
        sendMessageFromServerToClient(server, client);
    }

    @Test
    public void serverSendsEncryptedMessageToClientThatTrustsIt()
            throws IOException, InvalidInputException {
        ServerWithEncryption server = new ServerWithEncryption();
        Client client = new Client(server.getPublicSigningKey().getKeyId());
        server.setRecipient(client.getPublicEncryptionKey());

        sendMessageFromServerToClient(server, client);
    }

    @Test
    public void serverSendsEncryptedMessageToClientThatTrustsMasterServer()
            throws IOException, InvalidInputException {
        final Master master = new Master();
        final PartiallyTrustedServer server = new PartiallyTrustedServer();
        final Client client = new Client(master.getMasterPublicKeyId());
        server.setRecipient(client.getPublicEncryptionKey());
        server.setCertificate(master.delegateTrustTo(server
                .getPublicSigningKey()));

        sendMessageFromServerToClient(server, client);
    }

    private void sendMessageFromServerToClient(Server server, Client client)
            throws IOException, InvalidInputException, ParseException {
        final Service service = new Service("http", 80);
        final ByteOpenable message = server.generateMessage(service);
        final Service readBack = client.receiveMessage(message);
        assertThat(readBack.name, is(service.name));
        assertThat(readBack.port, is(service.port));
        PrettyPrinter.prettyPrint(new PrintWriter(System.out), message.read());
    }
}

@@ -117,16 +117,53 @@
 ertThat;
+%0Aimport static org.junit.Assert.fail;
 %0A%0Aimport
@@ -390,16 +390,70 @@
 enable;%0A
+import net.lshift.spki.suiteb.CryptographyException;%0A%0A
 import o
@@ -1954,24 +1954,552 @@
 nt);%0A    %7D%0A%0A
+    @Test%0A    public void untrustedServerSendsMessageToClient() throws IOException,%0A            InvalidInputException %7B%0A        Server trustedServer = new Server();%0A        Server attacker = new Server();%0A        Client client = new Client(trustedServer.getPublicSigningKey()%0A                .getKeyId());%0A        try %7B%0A            sendMessageFromServerToClient(attacker, client);%0A            fail(%22Expected client to fail to read any trusted content from message%22);%0A        %7D catch (CryptographyException e) %7B%0A        %7D%0A    %7D%0A%0A
     private 

Add test demonstrating untrusted server.