package ssh;

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

/**
 * This class implements the client side of the ssh protocol.
 * It establishes an ssh session on a channel, and then provides an
 * InputStream/OutputStream abstraction to allow the caller to transmit
 * data securely over the underlying channel.<p>
 * 
 * Source: draft-ylonen-ssh-protocol-00.txt<p>
 *
 * @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
 * 
 * @author Jon Howell <jonh@cs.dartmouth.edu> 
 * @rcs $Id: ClientProtocol.java,v 1.2 2000/05/22 01:35:50 jonh Exp $
 */
public class ClientProtocol {
	public static String clientVersion = "SSH-1.5-1.0.0\n";
		// *must* end in \n to comply with protocol

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

	public void connect(String host, Authenticator auth)
		throws IOException {
		connect(host, 22, auth);
	}

	public void connect(String host, int port, Authenticator auth)
		throws IOException {
		Authenticator[] auths = new Authenticator[1];
		auths[0] = auth;
		connect(host, port, auths);
	}

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

	public void connect(Socket socket, Authenticator[] auth)
		throws IOException {
		socketIn = socket.getInputStream();
		socketOut = socket.getOutputStream();
		this.auth = auth;

		exchangeVersionIdentification();
		binaryIn = new BinaryPacketInputStream(socketIn);
		binaryOut = new BinaryPacketOutputStream(socketOut);
		keyExchange();	// get server & host keys; reply with session key
						// (draft-ylonen page 8)
		checkSuccess();	// verify that server accepted session key
		sendUserName();	// identify incoming user; name only has to make
						// sense to server.
		authenticate();	// authenticate the identity of that user

		sshIn = new SshInputStream(binaryIn);
		sshOut = new SshOutputStream(binaryOut);
	}

	public void preparatory(boolean getPty)
		throws IOException {
		// Preparatory Operations phase (draft-ylonen page 13)
		// TODO: negotiate compression here

		if (getPty) {
			BinaryPacketOut bpo = binaryOut.newPacket();
			bpo.setType(Protocol.SSH_CMSG_REQUEST_PTY);
			bpo.writeString("vt100");
			bpo.writeInt(24);	// height, rows
			bpo.writeInt(80);	// width, columns
				// terminal shape should be determined by querying terminal,
				// but we can't in Java.
			bpo.writeInt(0);	// width, pixels
			bpo.writeInt(0);	// height, pixels
			bpo.writeByte(0);	// TTY_OP_END (end of terminal modes)
			binaryOut.writePacket(bpo);
			checkSuccess();
		}

		BinaryPacketOut bpo = binaryOut.newPacket();
		bpo.setType(Protocol.SSH_CMSG_EXEC_SHELL);
		binaryOut.writePacket(bpo);
	}

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

	public OutputStream getOutputStream() {
		return sshOut;
	}

	// an interface to directly run this class from the command line
	public static void main(String args[]) {
		try {
			String host = "localhost";
			int port = 22;
			if (args.length >= 1) {
				host = args[0];
			}
			if (args.length >= 2) {
				port = Integer.parseInt(args[1]);
			}
			System.out.println("Connecting to "+host+" on port "+port);
			Authenticator auth = new StdinPasswordAuthenticator();
			ClientProtocol cp = new ClientProtocol();
			cp.connect(host, port, auth);

			cp.preparatory(true);	// get a PTY

			Terminal t = new Terminal(cp.getInputStream(), cp.getOutputStream());
		} catch (Exception ex) {
			ex.printStackTrace();
		}
	}

	// instance variables
	InputStream socketIn;
	OutputStream socketOut;
	BinaryPacketInputStream binaryIn;
	BinaryPacketOutputStream binaryOut;
	SshInputStream sshIn;
	SshOutputStream sshOut;
	Authenticator[] auth;
	SshRandom random;		// source of random bits

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

	void exchangeVersionIdentification()
		throws IOException {
		// 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 constantSSH = tok.nextToken();
		String protocolVersion = tok.nextToken();
		String softwareVersion = tok.nextToken();
		if (!constantSSH.equals("SSH")) {
			throw new IOException("Remote version string is bad.");
		}

		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);
		}
		if (protMinor < 5) {
			throw new IOException("Remote minor version too old: "
				+protocolVersion);
		}

		// now send our version number to server
		socketOut.write(clientVersion.getBytes());
	}

	void keyExchange()
		throws IOException {

		// the first packet from the server is the public key
		BinaryPacketIn bpi = binaryIn.readPacket();
		assert(bpi.getType() == Protocol.SSH_SMSG_PUBLIC_KEY);
		if (bpi.getType() != Protocol.SSH_SMSG_PUBLIC_KEY) {
			throw new IOException("Unexpected packet type: "+bpi.getType()
				+"; expected "+Protocol.SSH_SMSG_PUBLIC_KEY);
		}
		byte[] cookie = new byte[8];
		bpi.readFully(cookie, 0, 8);
									// we return this; it helps prevent
									// IP spoofing in the rhosts case
		RSAKey serverKey				= RSAKey.readSsh(bpi);
		RSAKey hostKey					= RSAKey.readSsh(bpi);
		int protocolFlags				= bpi.readInt();
		int supportedCiphersMask		= bpi.readInt();
		int supportAuthenticationMask	= bpi.readInt();

		// TODO: should verify that we recognize the hostKey; otherwise
		// a trojan server could supply us with any old host key.

		// compute session ID
		byte[] sessionID = computeSessionID(hostKey, serverKey, cookie);

		// generate a random session key
		byte[] sessionKey = new byte[32];
		for (int i=0; i<32; i++) {
			sessionKey[i] = (byte) random.nextByte();
		}

		// the session key we return to the host is first xor'd with
		// the sessionID, then encrypted with the server and host keys.
		// (draft-ylonen, page 8)
		// [I don't really know what the value of the xor step is, since
		// the sessionID is easily snooped by an onlooker.]
		byte[] sessionKeyXorID = new byte[32];
		System.arraycopy(sessionKey, 0, sessionKeyXorID, 0, 32);
		for (int i=0; i<16; i++) {
			sessionKeyXorID[i] ^= sessionID[i];
		}
		// now public-key-encrypt, first with the smaller of the two
		// keys, then with the larger.
		RSAKey smallerKey, largerKey;
		if (serverKey.bits < hostKey.bits) {
			smallerKey = serverKey;
			largerKey = hostKey;
		} else {
			smallerKey = hostKey;
			largerKey = serverKey;
		}
		if (largerKey.bits - smallerKey.bits < 128) {	// draft-ylonen page 8
			throw new IOException(
				"Host and server keys must differ by 128 bits");
		}
		//BigInteger temp = smallerKey.encrypt(sessionKeyXorID, random);
		//BigInteger encryptedSessionKey = largerKey.encrypt(temp, random);
		// TODO: diddling to see if this works
		BigInteger temp = smallerKey.encrypt(sessionKeyXorID, random);
		BigInteger encryptedSessionKey =
			largerKey.encrypt(temp.toByteArray(), random);

		// send SSH_CMSG_SESSION_KEY packet. (draft-ylonen page 20)
		BinaryPacketOut bpo = binaryOut.newPacket();
		bpo.setType(Protocol.SSH_CMSG_SESSION_KEY);
		bpo.writeByte(Protocol.SSH_CIPHER_IDEA);
			// TODO: check that IDEA is actually supported
			// TODO: support other ciphers
		bpo.write(cookie);
		bpo.writeBigInteger(encryptedSessionKey);
		bpo.writeInt(0);		// no protocol extensions (?)
		binaryOut.writePacket(bpo);

		// Install cipher on input and output streams,
		// notably *after* sending session key to server.
		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);
	}

	void checkSuccess()
		throws IOException {
		// verify that the server accepted our SESSION_KEY packet
		BinaryPacketIn bpi = binaryIn.readPacket();
		if (bpi.getType() != Protocol.SSH_SMSG_SUCCESS) {
			throw new IOException("Unexpected packet type: "+bpi.getType()
				+"; expected "+Protocol.SSH_SMSG_SUCCESS);
		}
	}

	void sendUserName()
		throws IOException {
		
		BinaryPacketOut bpo = binaryOut.newPacket();
		bpo.setType(Protocol.SSH_CMSG_USER);
		bpo.writeString("jonh");
		binaryOut.writePacket(bpo);
	}

	public void authenticate()
		throws IOException {
		// read a packet; loop until SSH_SMSG_SUCCESS
		int authi=0;
		while (true) {
			BinaryPacketIn bpi = binaryIn.readPacket();
			System.out.println("result of auth: "+bpi.getType());
			if (bpi.getType() == Protocol.SSH_SMSG_SUCCESS) {
				break;
			} else if (bpi.getType() == Protocol.SSH_SMSG_FAILURE) {
				if (authi>=auth.length) {
					throw new IOException("all Authenticators failed.");
				}
				auth[authi].authenticate(binaryIn, binaryOut);
				authi++;
				continue;
			} else {
				throw new IOException(
					"Unexpected packet type in performAuthentication: "
					+bpi.getType());
			}
		}
		System.out.println("authenticated successfully!");
	}

	byte[] computeSessionID(RSAKey hostKey, RSAKey serverKey, byte[] cookie) {

		MD5 md5 = new MD5();

		// hash host key
		byte[] hostModulus = hostKey.modulus.toByteArray();
		int nonzeroBytes = hostKey.modulus.bitLength()/8;
		md5.update(hostModulus, hostModulus.length-nonzeroBytes,
			nonzeroBytes);	// skip leading zero bytes

		// hash server key
		byte[] serverModulus = serverKey.modulus.toByteArray();
		nonzeroBytes = serverKey.modulus.bitLength()/8;
		md5.update(serverModulus, serverModulus.length-nonzeroBytes,
			nonzeroBytes);	// skip leading zero bytes

		// hash cookie
		md5.update(cookie);

		byte[] sessionID = new byte[16];
		md5.getValue(sessionID);
		return sessionID;
	}

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