package de.hsadmin.core.util.dns;

import java.io.ByteArrayInputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.io.OutputStream;

public class DNSQuery {
	
	// DNS constants
	private final static int MAX_LABEL = 128;

	// DNS flag values
	private final static int QR_MASK = 0x8000;
	private final static int QR_SHIFT = 0;

	private final static int OPCODE_MASK = 0x7800;

	private final static int AA_MASK = 0x0400;
	private final static int TC_MASK = 0x0200;
	private final static int RD_MASK = 0x0100;
	private final static int RA_MASK = 0x0080;
	private final static int ZERO_MASK = 0x0070;
	private final static int RCODE_MASK = 0x000f;

	// DNS query/type constants
	public final static int TYPE_A = 1;
	public final static int TYPE_NS = 2;
	public final static int TYPE_MD = 3;
	public final static int TYPE_MF = 4;
	public final static int TYPE_CNAME = 5;
	public final static int TYPE_SOA = 6;
	public final static int TYPE_MB = 7;
	public final static int TYPE_MG = 8;
	public final static int TYPE_MR = 9;
	public final static int TYPE_NULL = 10;
	public final static int TYPE_WKS = 11;
	public final static int TYPE_PTR = 12;
	public final static int TYPE_HINFO = 13;
	public final static int TYPE_MINFO = 14;
	public final static int TYPE_MX = 15;
	public final static int TYPE_AXFR = 252;
	public final static int TYPE_ANY = 255;

	// DNS query classes
	public final static byte CLASS_INTERNET = 1;

	// query data members
	public int m_idQuery = 0x0101;
	
	// opcodes
	public final static int OP_NORM_QRY = 0x0000;
	public final static int OP_INVERSE_QRY = 0x0800;
	public final static int OP_SERVER_STAT = 0x1000;
	public int m_iOpCode = OP_NORM_QRY;

	// flags
	public boolean m_fQueryResult = false;
	public boolean m_fAuthAns = false;
	public boolean m_fTruncated = false;
	public boolean m_fRecurse = true;
	public boolean m_fRecursionAvail = false;
	
	private int m_iQryType = TYPE_A;
	private int m_iQryClass = CLASS_INTERNET;
	private String m_strQryName = null;

	// return codes
	public int m_iRCode = 0;
	public final static int RCODE_SUCCESS = 0;
	public final static int RCODE_NAME_ERROR = 0x0003;

	public int m_cQuestions = 1;
	public int m_cAnswerRRs = 0;
	public int m_cAuthRRs = 0;
	public int m_cInfoRRs = 0;

	DNSResourceRecord [] m_arrAns;
	DNSResourceRecord [] m_arrAuth;
	DNSResourceRecord [] m_arrInfo;

	public static String GetTypeDesc(int iType) {
		String str = null;
		
		switch (iType) {
		case TYPE_A: { str = new String("TYPE_A"); } break;
		case TYPE_NS: { str = new String("TYPE_NS"); } break;
		case TYPE_CNAME: { str = new String("TYPE_CNAME"); } break;
		case TYPE_SOA: { str = new String("TYPE_SOA"); } break;
		case TYPE_PTR: { str = new String("TYPE_PTR"); } break;
		case TYPE_HINFO: { str = new String("TYPE_HINFO"); } break;
		case TYPE_MX: { str = new String("TYPE_MX"); } break;
		case TYPE_AXFR: { str = new String("TYPE_AXFR"); } break;
		case TYPE_ANY: { str = new String("TYPE_ANY"); } break;
		default: { str = new String(Integer.toString(iType)); };
		}

		return str;
	}
	
	public DNSResourceRecord[] getDNSAnswerRecords() {
		return m_arrAns;
	}

	public DNSResourceRecord[] getDNSAuthorityRecords() {
		return m_arrAuth;
	}

	public static String GetClassDesc(int iClass) {
		String str = null;

		switch (iClass) {
		case CLASS_INTERNET: { str = new String("CLASS_INTERNET"); } break;
		default: { str = new String(); };
		}

		return str;
	}

	private int BuildFlags() {
		int iFlags = 0;

		iFlags |= m_iOpCode;

		if (m_fRecurse) {
			iFlags |= RD_MASK;
		}

		return iFlags;
	}

	private boolean SetFlags(int iFlags) {
		m_fQueryResult = (iFlags & QR_MASK) == QR_MASK;
		m_fAuthAns = (iFlags & AA_MASK) == AA_MASK;
		m_fTruncated = (iFlags & TC_MASK) == TC_MASK;
		m_fRecurse = (iFlags & RD_MASK) == RD_MASK;
		m_fRecursionAvail = (iFlags & RA_MASK) == RA_MASK;
		m_iRCode = iFlags & RCODE_MASK;

		return m_iRCode == RCODE_SUCCESS;
	}

	public boolean SetQuery(int iType, int iClass, String strName) {
		m_iQryType = iType;
		m_iQryClass = iClass;
		m_cQuestions = 1;
		m_strQryName = new String(strName);

		return true;
	}

	private boolean WriteFQDN(DataOutputStream dos, String strName) {
		try {
			int iPos = 0;
			int iSep = 0;

			while ((iSep = strName.indexOf('.', iPos)) >= 0) {
				dos.writeByte((byte)(iSep - iPos));
				dos.writeBytes(strName.substring(iPos, iSep));
				iPos = iSep + 1;
			}

			if (iPos < strName.length()) {
				dos.writeByte((byte)(strName.length() - iPos));
				dos.writeBytes(strName.substring(iPos));
			}
			
			// terminator
			dos.writeByte(0);

			return true;
		} catch (IOException ioe) {
			return false;
		}
	}
	
	public static String ReadLabelList(DataInputStream dis, byte [] abData) {
		int cb;
 		byte [] ab = new byte[MAX_LABEL];
		String str = new String();
		String strNext = new String();

		try {
			while ((cb = dis.readUnsignedByte()) > 0) {
				if ((cb & 0xc0) == 0xc0) {
					int iReadOff = (int)(cb & ~0xc0) << 8;

					cb = dis.readUnsignedByte();
					iReadOff += (int)cb;
					
					while (iReadOff < abData.length && ((cb = abData[iReadOff++]) > 0)) {
						if ((cb & 0xc0) == 0xc0) {
							return str;
						}

						strNext = new String(abData, 0, iReadOff, cb);

						if (str.length() > 0) {
							str += "." + strNext;
						} else {
							str = strNext;
						}

						iReadOff += cb;
					}

					return str;
				} else {
					byte [] abNext = new byte[cb];

					if (dis.read(abNext) < cb) {
						return new String();
					}
					
					strNext = new String(abNext, 0, 0, cb);
				}

				if (str.length() > 0) {
					str += "." + strNext;
				} else {
					str = strNext;
				}
			}
		} catch (IOException ioe) {
			System.err.println("exception: " + ioe.getMessage());
		}
	
		return str;
	}

	public boolean WriteQuery(OutputStream os) {
		try {
			if (m_strQryName == null) {
				return false;
			}

			DataOutputStream dos = new DataOutputStream(os);

			// query sequence ID
			dos.writeShort(m_idQuery);

			// query flags
			dos.writeShort(BuildFlags());

			// question count
			dos.writeShort(1);

			// answer count
			dos.writeShort(0);

			// authority count
			dos.writeShort(0);

			// additional info count
			dos.writeShort(0);

			// write query name
			WriteFQDN(dos, m_strQryName);

			// query type
			dos.writeShort(m_iQryType);

			// query class
			dos.writeShort(m_iQryClass);

			return true;
		} catch (IOException ioe) {	
			return false;
		}
	}

	private void SortRRs(DNSResourceRecord [] arr, boolean fDescending) {
		if (arr == null || arr.length < 2) {
			return;
		}

		boolean fSwapped;
		DNSResourceRecord rrSwap;

		do {
			fSwapped = false;
			
			for (int i = 0; i < arr.length - 1; i++) {
				boolean fSwap = false;

				if (arr[i+1].m_iType < arr[i].m_iType) {	
					fSwap = true;
				} else if (arr[i+1].m_iType == arr[i].m_iType) {
					switch (arr[i].m_iType) {
					case TYPE_MX: {
						fSwap = arr[i+1].m_lData < arr[i].m_lData;
					} break;
					}
				}

				if (fSwap = fSwap ^ fDescending) {
					rrSwap = arr[i];
					arr[i] = arr[i+1];
					arr[i+1] = rrSwap;
					fSwapped = true;
				}
			}
		} while (fSwapped);
	}

	public boolean ReadQuery(byte [] abData, int cbData) {
		try {
			ByteArrayInputStream is = new ByteArrayInputStream(abData, 0, cbData);

			DataInputStream dis = new DataInputStream(is);

			// query sequence ID
			m_idQuery = dis.readShort();

			// query flags
			if (!SetFlags(dis.readShort())) {
				return false;
			}

			// question count
			m_cQuestions = dis.readShort();

			// answer count
			m_cAnswerRRs = dis.readShort();

			// authority count
			m_cAuthRRs = dis.readShort();

			// additional info count
			m_cInfoRRs = dis.readShort();

			// read query name
			m_strQryName = ReadLabelList(dis, abData);

			// query type
			m_iQryType = dis.readShort();

			// query class
			m_iQryClass = dis.readShort();

			m_arrAns = new DNSResourceRecord[m_cAnswerRRs];

			for (int i = 0; i < m_cAnswerRRs; i++) {
				m_arrAns[i] = new DNSResourceRecord();

				if (!m_arrAns[i].readRecord(dis, abData)) {
					return false;
				}
			}

			SortRRs(m_arrAns, false);

			if (m_cAuthRRs > 0) {
				m_arrAuth = new DNSResourceRecord[m_cAuthRRs];

				for (int i = 0; i < m_cAuthRRs; i++) {
					m_arrAuth[i] = new DNSResourceRecord();

					if (!m_arrAuth[i].readRecord(dis, abData)) {
						return false;
					}
				}

				SortRRs(m_arrAuth, false);
			} else {
				m_arrAuth = null;
			}

			if (m_cInfoRRs > 0) {
				m_arrInfo = new DNSResourceRecord[m_cInfoRRs];

				for (int i = 0; i < m_cInfoRRs; i++) {
					m_arrInfo[i] = new DNSResourceRecord();

					if (!m_arrInfo[i].readRecord(dis, abData)) {
						return false;
					}
				}

				SortRRs(m_arrInfo, false);
			} else {
				m_arrInfo = null;
			}

		} catch (IOException ioe) {
			System.err.println("exception: " + ioe.getMessage());
			
			return false;
		}

		return true;
	}

}
