JmsWriter.java

package com.mcafee;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.Date;

import javax.jms.BytesMessage;
import javax.jms.JMSException;
import javax.jms.MapMessage;
import javax.jms.Message;
import javax.jms.MessageEOFException;
import javax.jms.ObjectMessage;
import javax.jms.StreamMessage;
import javax.jms.TextMessage;

import org.apache.commons.codec.binary.Hex;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.thoughtworks.xstream.XStream;

/**
 * Objects of this class write Message and its data to flat files. It determines
 * the type of message and then writes data as per different headers. Types of
 * messages supported are:<br/> 
 * 1. Message <br/> 
 * 2. TextMessage <br/>  
 * 3. ObjectMessage <br/> 
 * 4. BytesMessage<br/> 
 * 5. Stream Message <br/> 
 * 6. MapMessage<br/> 
 * 
 * @author Gursev Singh Kalra @ McAfee, Inc.
 * 
 */

public class JmsWriter {

	private static final Logger LOG = LoggerFactory.getLogger(JmsWriter.class);
	
	private static String fileSeparator;
	private String targetDir;
	private boolean targetDirExists;
	private static String msgSeparator; // used as a separator between messages
										// being dumped
	private String msgDumpFilePath;
	private PrintWriter msgDumpFileHandle;	

	private XStream jmsXstream;	

	private int textMsgsWritten = 0;
	private int bytesMsgsWritten = 0;
	private int objMsgsWritten = 0;
	private int streamMsgsWritten = 0;
	private int mapMsgsWritten = 0;
	private int msgsWritten = 0;

	private String identifier;

	static {
		msgSeparator = "\n~-~-~-~-~-~-~-~-~-~~-~-~-~-~-~-~-~-~-~" + (new Date()).getTime() + "\n";
	}

	public JmsWriter(String targetDir, String identifier) throws JmsDiggerException {
		this.targetDir = targetDir;
		targetDirExists = false;
		fileSeparator = System.getProperty("file.separator");
		try {
			(new File(targetDir)).mkdirs();
		} catch (Exception ex) {
			LOG.error("Target directory could not be created", ex);
			throw JmsHelper.buildJmsDiggerException("Target directory could not be created");
		}
		jmsXstream = new XStream();		
		this.identifier = identifier;
	}
	
	public void init() throws FileNotFoundException, IOException {
		LOG.debug("Entering init");
		initFileHandles();
		LOG.debug("Leaving init");
	}

	private void initFileHandles() throws FileNotFoundException, IOException {
		// time in seconds since epoch
		LOG.debug("Entering initFileHandles");
		long dt = ((new Date()).getTime());
		msgDumpFilePath = targetDir + fileSeparator + identifier + "-MessageDump-" + dt + ".txt";

		try {
			msgDumpFileHandle = new PrintWriter(new FileOutputStream(msgDumpFilePath), true);
			msgDumpFileHandle.println("#########################################################################\nJMS Destination message dump generated by JMSDigger\n#########################################################################\n");
		} catch (FileNotFoundException ex) {
			LOG.error("Output file handle could not be creatd", ex);
			throw ex;
		}
		
		LOG.debug("Leaving initFileHandles");
	}

	private String byteArrayToHexString(byte[] b) {
		LOG.debug("Entering byteArrayToHexString");
		// Converts a byte array to string representation of hex digits
		// { 0x12, 0x23, 0x32, 0xA5 }
		StringBuilder sb = new StringBuilder();
		if (b.length == 0)
			return "{ }";

		String str = Hex.encodeHexString(b);
		String[] twoCharArray = str.split("(?<=\\G.{2})");

		sb.append("{ ");
		for (String s : twoCharArray)
			sb.append("0x" + s + ", ");

		sb.deleteCharAt(sb.length() - 2);
		sb.append("}");
		LOG.debug("Leaving byteArrayToHexString");
		return sb.toString();
	}

	private String getBytesMessageContentAsString(BytesMessage msg) throws JmsDiggerException {
		LOG.debug("Entering writeBytesMessage");
		int len = -1;
		byte[] body;

		try {
			len = (int) msg.getBodyLength();
			if (len == 0) {
				return "{ }";
			}
	
			body = new byte[len];
			msg.readBytes(body);
		} catch (JMSException ex) {
			LOG.info("An error occured while retreving content from ByteMessage", ex);
			throw JmsHelper.buildJmsDiggerException("An error occured while retreving content from ByteMessage", ex);
		}
		LOG.debug("Leaving writeBytesMessage");
		return byteArrayToHexString(body);

	}

	private String getStreamMessageContentAsString(StreamMessage msg) throws JmsDiggerException {
		LOG.debug("Entering writeStreamMessage");
		StringBuilder sb = new StringBuilder();
		sb.append("{ ");
		Object o;
		
		try {
			while (true) {
				o = msg.readObject();
				if (o instanceof byte[])
					sb.append(byteArrayToHexString((byte[]) o));
				else
					sb.append(o);
				sb.append(", ");
			}
		} catch (MessageEOFException ex) {
			sb.deleteCharAt(sb.length() - 2);
			// Do nothing. End of StreamMessage content has reached.
		} catch (JMSException ex) {
			LOG.info("An error occured while reading Object from StreamMessage", ex);
			throw JmsHelper.buildJmsDiggerException("An error occured while reading Object from StreamMessage", ex);
		}
		
		sb.append("}");
		LOG.debug("Leaving writeStreamMessage");
		return sb.toString();		
	}

	private String getMapMessageContentAsString(MapMessage msg) throws JmsDiggerException {
		LOG.debug("Entering writeMapMessage");
		try {			
			return(JmsHelper.mapMessageToString(msg, "[+] MapMessage name-value pairs : "));
		} catch (JMSException ex) {
			LOG.info("An error occured while reading MapMessage", ex);
			throw JmsHelper.buildJmsDiggerException("An error occured while reading MapMessage", ex);
		}
	}


	/**
	 * Analyze the message type against all the JMS message types supported and
	 * write metadata as per the message type
	 * 
	 * @param msg
	 * @param result
	 * @throws JmsDiggerException 
	 * @throws IOException
	 */
	public void writeMsg(Message msg) throws JmsDiggerException {
		LOG.debug("Entering writeMsg");
		if (msg == null)
			throw new IllegalArgumentException("Message parameter is null");

		JmsMsgType type = null;
		StringBuilder messageContent = new StringBuilder();

		try {
			if (msg instanceof BytesMessage) {
				LOG.debug("Message Type : BytesMessage");
				type = JmsMsgType.BYTES;
				messageContent.append("[+] Message Type : BytesMessage");
				messageContent.append("\n");
				messageContent.append("[+] Message Bytes as an Array: \n");
				messageContent.append(getBytesMessageContentAsString((BytesMessage) msg));
				messageContent.append("\n");
				bytesMsgsWritten++;
			}

			if (type == null && msg instanceof TextMessage) {
				LOG.debug("Message Type : TextMessage");
				type = JmsMsgType.TEXT;
				messageContent.append("[+] Message Type : TextMessage");
				messageContent.append("\n");
				messageContent.append("[+] Extracted Text\n");
				messageContent.append(((TextMessage)msg).getText());
				messageContent.append("\n");
				textMsgsWritten++;
			}

			if (type == null && msg instanceof StreamMessage) {
				LOG.debug("Message Type : StreamMessage");
				type = JmsMsgType.STREAM;
				messageContent.append("[+] Message Type : StreamMessage");
				messageContent.append("\n");
				messageContent.append("[+] Extracted Stream\n");
				messageContent.append(getStreamMessageContentAsString((StreamMessage) msg));
				messageContent.append("\n");
				streamMsgsWritten++;
			}

			if (type == null && msg instanceof MapMessage) {
				LOG.debug("Message Type : MapMessage");
				type = JmsMsgType.MAP;
				messageContent.append("[+] Message Type : MapMessage");
				messageContent.append("\n");
				messageContent.append(getMapMessageContentAsString((MapMessage) msg));
				messageContent.append("\n");
				mapMsgsWritten++;
			}

			if (type == null && msg instanceof ObjectMessage) {
				LOG.debug("Message Type : ObjectMessage");
				type = JmsMsgType.OBJECT;
				messageContent.append("[+] Message Type : ObjectMessage");
				messageContent.append("\n");
				objMsgsWritten++;
			}

			if (type == null && msg instanceof Message) {
				LOG.debug("Message Type : Message");
				type = JmsMsgType.MESSAGE;
				messageContent.append("[+] Message Type : Message");
				messageContent.append("\n");
			}

			messageContent.append("[+] Message Properties in XML (including binary dump in content/data element)\n");
			messageContent.append(jmsXstream.toXML(msg));
			messageContent.append("\n");
			messageContent.append(msgSeparator);
			msgDumpFileHandle.println(messageContent.toString());
			msgsWritten++;
		} catch (JMSException ex) {
			LOG.info("A JMS operation failed", ex);
			throw JmsHelper.buildJmsDiggerException("A JMS operation failed", ex);
		}
		LOG.debug("Leaving writeMsg");
	}

	private String getWriteCountBreakup() {
		StringBuilder sb = new StringBuilder();
		sb.append("ObjectMessages written: " + objMsgsWritten + ", ");
		sb.append("BytesMessages written: " + bytesMsgsWritten + ", ");
		sb.append("StreamMessages written: " + streamMsgsWritten + ", ");
		sb.append("MapMessages written: " + mapMsgsWritten + ", ");
		sb.append("TextMessages written: " + textMsgsWritten);		
		return sb.toString();
	}

	public void close() throws IOException {
		LOG.debug("Entering close");
		if (msgsWritten == 0)
			msgDumpFileHandle.println("No messages were written");

		msgDumpFileHandle.flush();
		msgDumpFileHandle.close();
		
		LOG.debug(getWriteCountBreakup());
		LOG.debug("Leaving close");
	}

}