package ssh;

import ssh.RSA.*;
import Tools.*;
import java.io.*;
import java.net.*;
import java.util.*;
import java.math.*;

/**
 * This class implements a simple encrypted channel, based on
 * ssh. Several features are left out (man-in-the-middle and privacy
 * defenses), meant to be implemented and verified by a higher layer
 * based on the regarding calculus.<p>
 *
 * Source: based on ClientProtocol.java, my Java implementation of ssh
 * 1.5.
 *
 * @deprecated This is the version of the protocol that predates
 * JDK1.2's RMISocketFactory mechanism. See {@link SSHServerSocketFactory}
 * for the new way to plug this SSH channel implementation into RMI.<p>
 * 
 * @author Jon Howell <jonh@cs.dartmouth.edu>
 * 
 * @license This code is Copyright 1999 Jon Howell. It is available for
 * use under the GNU Public License, available at:
 * http://www.gnu.org/copyleft/gpl.html
 * 
 * @rcs $Id: SecureRMIProtocol.java,v 1.4 2000/05/26 06:17:34 jonh Exp $
 */

public class SecureRMIProtocol
	implements SRPConstants {

	public static String protocolVersion = "SRP-1.0-1.0.0\n";
		// Protocol string is like ssh protocol string.
		// 'SRP' identifies the protocol,
		// 1.0 identifies the protocol version,
		// 1.0.0 identifies the client software version.
		// (sort of useless if there are multiple implementations.)

	public SecureRMIProtocol() {
		random = new SshRandom();
		// TODO: use some more meaningful random seed
		// wait, make that *less* meaningful. :v)
		sshIn = null;
		sshOut = null;
		oppositeKey = null;
	}

	public void setKey(RSAKey[] pair) {
		setKey(pair[0], pair[1]);
	}

	public void setKey(RSAKey privateKey, RSAKey publicKey) {
		this.privateKey = privateKey;
		this.publicKey = publicKey;
	}

	public RSAKey getOppositeKey() {
		return oppositeKey;
	}

	public void connect(String host, int port)
		throws IOException {
		connect(new Socket(host, port));
	}

	public void connect(Socket socket)
		throws IOException {
		connect(socket.getInputStream(), socket.getOutputStream());
	}

	public void connect(InputStream is, OutputStream os)
		throws IOException {
		socketIn = is;
		socketOut = os;

		log.setPrefix("connect: ").log("verbose", "exchangeVersion");
		exchangeVersionIdentification();
		binaryIn = new BinaryPacketInputStream(socketIn);
		binaryOut = new BinaryPacketOutputStream(socketOut);
		log.log("verbose", "key exchange");
		clientKeyExchange();	// prove we each have our respective keys;
								// establish session key

		// these are the streams a higher-level protocol will use to
		// actually transmit data.
		sshIn = new SshInputStream(binaryIn);
		sshOut = new SshOutputStream(binaryOut);
		log.log("verbose", "done");
	}

	public void accept(Socket socket)
		throws IOException {
		accept(socket.getInputStream(), socket.getOutputStream());
	}

	public void accept(InputStream is, OutputStream os)
		throws IOException {
		socketIn = is;
		socketOut = os;

		log.setPrefix("accept: ").log("verbose", "exchangeVersion");
		exchangeVersionIdentification();
		binaryIn = new BinaryPacketInputStream(socketIn);
		binaryOut = new BinaryPacketOutputStream(socketOut);
		log.log("verbose", "key exchange");
		serverKeyExchange();	// prove we each have our respective keys;
								// establish session key

		// these are the streams a higher-level protocol will use to
		// actually transmit data.
		sshIn = new SshInputStream(binaryIn);
		sshOut = new SshOutputStream(binaryOut);
		log.log("verbose", "done");
	}

	// these become meaningful once connect has returned successfully
	public InputStream getInputStream() {
		return sshIn;
	}

	public OutputStream getOutputStream() {
		return sshOut;
	}

	// instance variables
	InputStream socketIn;
	OutputStream socketOut;
	BinaryPacketInputStream binaryIn;
	BinaryPacketOutputStream binaryOut;
	SshInputStream sshIn;
	SshOutputStream sshOut;
	SshRandom random;		// source of random bits
	RSAKey privateKey;		// keys for local end of connection
	RSAKey publicKey;
	RSAKey oppositeKey;		// key received for remote end of connection

	static Log log = new Log();
	static {
		// log.addLevel("verbose");
		// log.addLevel("packets");
	}

	////////////////////////////////////////////////////////////////////////
	//  private methods  ///////////////////////////////////////////////////
	////////////////////////////////////////////////////////////////////////

	void exchangeVersionIdentification()
		throws IOException {

		// send our version number to server
		socketOut.write(protocolVersion.getBytes());

		// read characters until a \r or \n. Don't read past that!
		// after that, the protocol switches to ssh's binary packet protocol
		StringBuffer sb = new StringBuffer();
		while (true) {
			int ci = socketIn.read();
			if (ci<0) throw new EOFException();
			char c = (char) ci;
			if (c == '\r') {
				c = '\n';
			}
			sb.append(c);
			if (c == '\n') {
				break;
			}
		}

		String remoteVersion = sb.toString();
		StringTokenizer tok = new StringTokenizer(remoteVersion, "-");
		String constantSRP = tok.nextToken();
		String protocolVersion = tok.nextToken();
		String softwareVersion = tok.nextToken();
		if (!constantSRP.equals("SRP")) {
			throw new IOException("Remote version string is bad: "
				+remoteVersion);
		}

		StringTokenizer tok2 = new StringTokenizer(protocolVersion, ".");
		int protMajor = Integer.parseInt(tok2.nextToken());
		int protMinor = Integer.parseInt(tok2.nextToken());
		if (protMajor != 1) {
			throw new IOException("Remote major version incompatible: "
				+protocolVersion);
		}
	}

	void clientKeyExchange()
		throws IOException {

		// the first packet from the server is the server key
		BinaryPacketIn bpi = binaryIn.readPacket();
		assert(bpi.getType() == SRP_SMSG_SERVER_KEY);
		if (bpi.getType() != SRP_SMSG_SERVER_KEY) {
			throw new IOException("Unexpected packet type: "+bpi.getType()
				+"; expected "+SRP_SMSG_SERVER_KEY);
		}
		RSAKey serverKey				= RSAKey.readSsh(bpi);
		int protocolFlags				= bpi.readInt();
		int supportedCiphersMask		= bpi.readInt();
		BigInteger clientChallenge		= StreamExtras.readBigInteger(bpi);

		// take the challenge (prove we hold the private side of
		// the client key we're sending; prevents replay attack)
		BigInteger eClientChallenge
			= privateKey.encrypt(clientChallenge.toByteArray(), random);

		// create our own challenge for the server
		BigInteger serverChallenge = random.newBigInteger(32);

		// send SRP_CMSG_CLIENT_KEY packet.
		BinaryPacketOut bpo = binaryOut.newPacket();
		bpo.setType(SRP_CMSG_CLIENT_KEY);
		publicKey.writeSsh(bpo);
		StreamExtras.writeBigInteger(bpo, eClientChallenge);
		StreamExtras.writeBigInteger(bpo, serverChallenge);
		bpo.writeByte(SRP_CIPHER_IDEA);
		binaryOut.writePacket(bpo);

		// receive SRP_SMSG_SESSION_KEY packet
		bpi = binaryIn.readPacket();
		assert(bpi.getType() == SRP_SMSG_SESSION_KEY);
		BigInteger eServerChallenge		= StreamExtras.readBigInteger(bpi);
		BigInteger eSessionKey			= StreamExtras.readBigInteger(bpi);

		// verify server challenge was met
		byte[] dec = serverKey.decrypt(eServerChallenge);
		if (dec == null) {
			throw new IOException(
				"server's response does not decrypt to challenge");
		}
		BigInteger serverChallengeReply = new BigInteger(dec);
		if (!serverChallengeReply.equals(serverChallenge)) {
			throw new IOException(
				"server's response does not decrypt to challenge");
		}

		oppositeKey = serverKey;

		// decrypt session key (for my eyes only!)
		BigInteger sessionKeyBI =
			new BigInteger(privateKey.decrypt(eSessionKey));
		byte[] sessionKey = sessionKeyBI.toByteArray();

		Cipher cipher = new CipherIdea();
		cipher.setKey(sessionKey);
		binaryOut.setCipher(cipher);

		// need separate cipher objects on each stream, even though
		// the key is the same, because streams are chained independently.
		// (draft-ylonen page 5). [Makes sense, since packets are sent
		// asynchronously from each end.]
		cipher = new CipherIdea();
		cipher.setKey(sessionKey);
		binaryIn.setCipher(cipher);

		// send success message to indicate (indeed, demonstrate)
		// that the session key is good.
		bpo = binaryOut.newPacket();
		bpo.setType(SRP_CMSG_SUCCESS);
		binaryOut.writePacket(bpo);
	}

	void serverKeyExchange()
		throws IOException {

		// create a challenge for the client
		BigInteger clientChallenge = random.newBigInteger(32);

		// the first packet from the server is the server key
		BinaryPacketOut bpo = binaryOut.newPacket();
		bpo.setType(SRP_SMSG_SERVER_KEY);
		publicKey.writeSsh(bpo);
		bpo.writeInt(0);						// protocolFlags
		bpo.writeInt(SRP_CIPHER_IDEA);			// supportedCipherMask
		StreamExtras.writeBigInteger(bpo, clientChallenge);
		binaryOut.writePacket(bpo);

		// receive client's response and challenge for us
		BinaryPacketIn bpi = binaryIn.readPacket();
		assert(bpi.getType() == SRP_CMSG_CLIENT_KEY);
		RSAKey clientKey				= RSAKey.readSsh(bpi);
		BigInteger eClientChallenge		= StreamExtras.readBigInteger(bpi);
		BigInteger serverChallenge		= StreamExtras.readBigInteger(bpi);
		byte selectedCipher				= bpi.readByte();

		// verify clientChallenge was responded to successfully
		byte[] dec = clientKey.decrypt(eClientChallenge);
		if (dec == null) {
			throw new IOException(
				"client's response does not decrypt to challenge");
		}
		BigInteger clientChallengeReply = new BigInteger(dec);
		if (!clientChallengeReply.equals(clientChallenge)) {
			throw new IOException(
				"client's response does not decrypt to challenge");
		}

		oppositeKey = clientKey;

		// ensure cipher selection is okay
		assert(selectedCipher == SRP_CIPHER_IDEA);

		// take the server challenge (prove we hold the private side of
		// the server key we're sending; prevents replay attack)
		BigInteger eServerChallenge
			= privateKey.encrypt(serverChallenge.toByteArray(), random);

		// create session key
		byte[] sessionKey
			= random.newByteArray(32);	// 32 bytes of randomness
		BigInteger eSessionKey = clientKey.encrypt(sessionKey, random);

		// send session key packet.
		bpo = binaryOut.newPacket();
		bpo.setType(SRP_SMSG_SESSION_KEY);
		StreamExtras.writeBigInteger(bpo, eServerChallenge);
		StreamExtras.writeBigInteger(bpo, eSessionKey);
		binaryOut.writePacket(bpo);

		// need separate cipher objects on each stream, even though
		// the key is the same, because streams are chained independently.
		// (draft-ylonen page 5). [Makes sense, since packets are sent
		// asynchronously from each end.]
		Cipher cipher = new CipherIdea();
		cipher.setKey(sessionKey);
		binaryOut.setCipher(cipher);

		cipher = new CipherIdea();
		cipher.setKey(sessionKey);
		binaryIn.setCipher(cipher);

		// receive SRP_CMSG_SUCCESS packet
		bpi = binaryIn.readPacket();
		assert(bpi.getType() == SRP_CMSG_SUCCESS);
	}

	private void assert(boolean condition)
		throws IOException {
		if (!condition) {
			throw new IOException("assertion failed");
		}
	}
}
