package dcj.examples.security;

import java.lang.*;
import java.net.*;
import java.io.*;
import java.security.*;

/**
 * Source code from "Java Distributed Computing", by Jim Farley.
 *
 * Class: SecureAgent
 * Example: 5-5
 * Description: An agent that includes the digital signature authentication
 *      of the AuthAgent, and adds the ability to encrypt messages using a
 *      symmetric cipher.
 * NOTE: This class requires the JCE package, available only in the US
 *      and Canada.
 */

public class SecureAgent extends SimpleAgent {

  // The Identity of the agent we're connected to
  Identity remoteAgent = null;
  // A secret key used to encode/decode messages
  Key cryptKey = null;

  public SecureAgent(String host, int port)
      throws IllegalArgumentException {

    super(host, port);
    DataInputStream din = new DataInputStream(inStream);

    // Try to authenticate the remote agent
    try {
      String agentId = din.readUTF();
      int dataLen = din.readInt();
      byte[] data = new byte[dataLen];
      din.read(data);
      int sigLen = din.readInt();
      byte[] sig = new byte[sigLen];
      din.read(sig);

      if (!authenticate(agentId, data, sig)) {
        // Failed to authenticate, write error message, close socket and
        // return
        System.out.println("Failed to authenticate remote agent "
                           + agentId);
        closeConnection();
      }
      else {
        // Remote agent is authenticated, first message is a welcome
        addMsg("HELLO " + agentId);
      }
    }
    catch (Exception e) {
      closeConnection();
    }
  }

  protected boolean authenticate(String id,
                                 byte[] data, byte[] sig) {
    boolean success = false;
    PublicKey key = lookupKey(id);
    try {
      // Set up a signature with the agent's public key
      Signature agentSig = Signature.getInstance(key.getAlgorithm());
      agentSig.initVerify(key);
      // Try to verify the signature message from the agent
      agentSig.update(data);
      success = agentSig.verify(sig);

      if (success) {
        // Agent checks out, so initialize an identity for it
        remoteAgent = new Identity(id);
        remoteAgent.setPublicKey(key);
        // Get the agent's secret encryption key, too
        cryptKey = lookupSecretKey(id);
      }
    }
    catch (Exception e) {
      System.err.println("Failed to verify agent signature.");
      success = false;
    }

    return success;
  }

  public void run() {
    // Go into infinite loop, sending messages, receiving responses and
    // processing them...

    DataInputStream din = (DataInputStream)inStream;
    DataOutputStream dout = (DataOutputStream)outStream;

    // Make an encryption Cipher for sending messages...
    String cryptAlgName = cryptKey.getAlgorithm();
    Cipher sendCipher = Cipher.getInstance(cryptAlgName);
    sendCipher.initEncrypt(cryptKey);
    // ...and a decryption Cipher for receiving them.
    Cipher receiveCipher = Cipher.getInstance(cryptAlgName);
    receiveCipher.initDecrypt(cryptKey);

    while (true) {
      String msg = nextMsg();
      if (msg != null) {
        String inMsg = "", inToken = "";
        try {
          // Send encrypted message to agent
          byte[] eData = sendCipher.crypt(msg.getBytes());
          dout.write(eData);

          // Read and decrypt message from agent
          int dataLen = din.readInt();
          eData = new byte[dataLen];
          din.read(eData);
          byte[] clearData = receiveCipher.crypt(eData);
          inMsg = new String(clearData);

          // Process the incoming message
          processMsg(inMsg);
        }
        catch (Exception e) {}
      }
    }
  }
}