CrossJeeves/src/com/flaremicro/crossjeeves/net/ClientHandler.java

package com.flaremicro.crossjeeves.net;

import static com.flaremicro.crossjeeves.net.ErrorCodes.*;

import java.io.BufferedInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.net.InetAddress;
import java.net.Socket;

import com.flaremicro.crossjeeves.net.packet.Packet;
import com.flaremicro.crossjeeves.net.packet.Packet0Identify;
import com.flaremicro.crossjeeves.net.packet.Packet127KeepAlive;
import com.flaremicro.crossjeeves.net.packet.Packet1Status;
import com.flaremicro.crossjeeves.net.packet.Packet2Script;
import com.flaremicro.crossjeeves.net.packet.Packet3Clone;
import com.flaremicro.crossjeeves.net.packet.Packet4FileData;
import com.flaremicro.crossjeeves.net.packet.Packet5Artifact;
import com.flaremicro.crossjeeves.net.packet.Packet6Disconnect;
import com.flaremicro.util.Util;
import com.flaremicro.util.ZipUtils;

public class ClientHandler extends NetworkHandler {
	private String script;
	
	public ClientHandler(Socket socket, String script) throws IOException {
		super(socket);
		this.script = script;
	}
	
	private void init() {
		enqueue(new Packet0Identify(0));
		this.beginWriteThread();
		this.beginReading();
	}
	
	@Override
	public void handlePacket(Packet packet) {
		disconnect(INVALID_PACKET_RECIEVED.id, "Recieved invalid packet " + packet.getId() + " (Unknown)");
	}

	@Override
	public void handlePacket(Packet0Identify packet) {
		if (packet.getProtocolVersion() != Packet.PROTOCOL_VERSION)
		{
			disconnect(OUTDATED_AGENT.id, "Recieved outdated protocol version " + packet.getProtocolVersion() + " (Expected " + Packet.PROTOCOL_VERSION + ")");
		}
		else
		{
			enqueue(new Packet2Script(script));
		}
	}

	@Override
	public void handlePacket(Packet1Status packet) {
		if((packet.getFlags() & Packet1Status.BUSY) != 0)
		{
			disconnect(OUTDATED_AGENT.id, "Agent is too busy");
		}
	}

	@Override
	public void handlePacket(Packet2Script packet) {
		disconnect(INVALID_PACKET_RECIEVED.id, "Recieved invalid packet " + packet.getId() + " (Script) (Server should not send a script)");
	}

	@Override
	public void handlePacket(Packet3Clone packet) {
		long fileId = packet.getFileID();
		String workspace = System.getenv("WORKSPACE");
		if (workspace == null || workspace.trim().length() <= 0)
		{
			disconnect(INVALID_SYSTEM_STATE.id, "Workspace is not defined");
		}
		BufferedInputStream fileStream = null;
		try
		{
			File workspaceZip = File.createTempFile("workspace-" + fileId, ".zip");
			ZipUtils.zipDirectory(new File(workspace), workspaceZip);
			packet = new Packet3Clone(fileId, workspaceZip.length());
			enqueue(packet);
			fileStream = new BufferedInputStream(new FileInputStream(workspaceZip));
			int read;
			byte[] buffer = new byte[4096];
			while ((read = fileStream.read(buffer)) > -1)
			{
				if (read == 0)
					continue;
				Packet4FileData dataPacket = new Packet4FileData(fileId, (short) read, buffer);
				enqueue(dataPacket);
			}
			enqueue(new Packet4FileData(fileId, (short) 0, new byte[] {}));
		}
		catch (IOException e)
		{
			disconnect(FILE_UPLOAD_FAILURE.id, "Failed to upload file");
			e.printStackTrace();
		}
		finally
		{
			Util.cleanClose(fileStream);
		}
	}

	@Override
	public void handlePacket(Packet4FileData packet) {
		try
		{
			this.appendFile(packet.getFileId(), packet.getFileChunk());
		}
		catch (IOException e)
		{
			e.printStackTrace();
			disconnect(FILE_DOWNLOAD_FAILURE.id, "Failed to download transferred file");
		}
	}

	@Override
	public void handlePacket(Packet5Artifact packet) {

	}

	@Override
	public void handlePacket(Packet127KeepAlive packet) {
		enqueue(packet);
	}

	@Override
	public void handlePacket(Packet6Disconnect packet) {
		System.out.printf("Disconnect %d: %s%n", packet.getCode(), packet.getReason());
		doDisconnect(packet.getCode());
	}

	public static int connect(InetAddress addr, int port, String script) {
		try
		{
			Socket socket = new Socket(addr, port);
			ClientHandler clientHandler = new ClientHandler(socket, script);
			clientHandler.init();
			return clientHandler.getExitCode();
		}
		catch (IOException e)
		{
			e.printStackTrace();
			return CONNECT_FAILED.id;
		}
	}

}