hostage/src/de/tudarmstadt/informatik/hostage/protocol/SMBUtils/NBNS.java

package de.tudarmstadt.informatik.hostage.protocol.SMBUtils;

import java.nio.ByteBuffer;

import de.tudarmstadt.informatik.hostage.commons.HelperUtils;


public class NBNS {
	
	private byte[] transactID;
	private byte[] flags;
	private byte[] questions;
	private byte[] answerRRs;
	private byte[] authorityRRs;
	private byte[] additionalRRs;
	private byte[] payload;
	private byte[] additional;
	private byte[] addr;
	private byte[] name;
	private int type;
	private int service;
	
	public NBNS(byte[] packet) {
		transactID = new byte[]{packet[0], packet[1]};
		flags = new byte[]{packet[2], packet[3]};
		questions = new byte[]{packet[4], packet[5]};
		answerRRs = new byte[]{packet[6], packet[7]};
		authorityRRs = new byte[]{packet[8], packet[9]};
		additionalRRs = new byte[]{packet[10], packet[11]};
		int length = 0;
		for (int i = 12; i < packet.length; i++, length++) {
			if (packet[i] == 0x01) {
				length++;
				break;
			}
		}
		payload = new byte[length];
		for (int i = 0; i < payload.length; i++) {
			payload[i] = packet[i+12];
		}
		additional = new byte[packet.length-12-length];
		for (int i = 0; i < additional.length; i++) {
			additional[i] = packet[i+12+length];
		}
	}
	
	public NBNS(byte[] transactID, String addr) {
		this.transactID = transactID;
		addressToBytes(addr);
	}
	
	public NBNS(byte[] transactID, int type, int service, String name, String addr) {
		this(transactID, addr);
		this.type = type;
		this.service = service;
		setName(name);
	}
	
	private void preparePacket() {
		switch (type) {
		case Type.REGISTRATION_UNIQUE:
			prepareRegistrationPacket();
			break;
		case Type.REGISTRATION_GROUP:
			prepareRegistrationPacket();
			break;
		case Type.NAME_QUERY:
			prepareNameQueryPacket();
			break;
		default:
		}
	}
	
	private void prepareRegistrationPacket() {
		flags = new byte[]{0x29, 0x10};
		questions = new byte[]{0x00, 0x01};
		answerRRs = new byte[]{0x00, 0x00};
		authorityRRs = new byte[]{0x00, 0x00};
		additionalRRs = new byte[]{0x00, 0x01};
		payload = getPayload();
		additional = getAdditionalRecords();
	}
	
	private void prepareNameQueryPacket() {
		flags = new byte[] {0x01, 0x10};
		questions = new byte[]{0x00, 0x01};
		answerRRs = new byte[]{0x00, 0x00};
		authorityRRs = new byte[]{0x00, 0x00};
		additionalRRs = new byte[]{0x00, 0x00};
		payload = getPayload();
	}
	
	private byte[] getPayload() {
		byte[] payload = NMBStringCoder.wrapNBNSName(this.name, service);
		byte[] type = {0x00, 0x20};
		byte[] nbnsclass = {0x00, 0x01};
		return HelperUtils.concat(payload, type, nbnsclass);
	}
		
	private void addressToBytes(String addrString) {
		String[] addrParts = addrString.split("\\.");
		addr = new byte[4];
		addr[0] = (byte)Integer.parseInt(addrParts[0]);
		addr[1] = (byte)Integer.parseInt(addrParts[1]);
		addr[2] = (byte)Integer.parseInt(addrParts[2]);
		addr[3] = (byte)Integer.parseInt(addrParts[3]);
	}
	
	public static byte[] getServiceBytes(int service) {
		switch (service) {
		case Service.SERVER:
			return new byte[]{0x43, 0x41};
		case Service.MESSENGER:
			return new byte[]{0x41, 0x44};
		case Service.WORKSTATION:
			return new byte[]{0x41, 0x41};
		case Service.BROWSER_ELECTION:
			return new byte[]{0x42, 0x4f};
		case Service.LOCAL_MASTER_BROWSER:
			return new byte[]{0x42, 0x4e};
		default:
			return new byte[]{0x43, 0x41};
		}
	}

	private byte[] getAdditionalRecords() {
		byte[] name = {(byte) 0xc0, 0x0c};
		byte[] type = {0x00, 0x20};
		byte[] nbnsclass = {0x00, 0x01};
		byte[] timeToLive = {0x00, 0x00, 0x00, 0x00};
		byte[] nameFlags = ((this.type == Type.REGISTRATION_UNIQUE) ? new byte[]{0x00, 0x00} : new byte[]{0x08, 0x00});
		byte[] buffer = ByteBuffer.allocate(4).putInt(nameFlags.length + addr.length).array();
		byte[] length = {buffer[2], buffer[3]};
		return HelperUtils.concat(name, type, nbnsclass, timeToLive, length, nameFlags, addr);
	}
	
	public byte[] getNextPacket() {
		preparePacket();
		transactID[1] += 0x01;
		return getBytes();
	}

	public byte[] getBytes() {
		return HelperUtils.concat(transactID, flags, questions, answerRRs, authorityRRs, additionalRRs, payload, additional);
	}

	public void setName(String name) {
		this.name = name.getBytes();
		this.name = NMBStringCoder.encodeNBNSName(this.name);
	}

	public void setType(int type) {
		this.type = type;
	}

	public void setService(int service) {
		this.service = service;
	}

	public byte[] getAndIncrementTransactID() {
		byte[] response = new byte[2];
		response[0] = transactID[0];
		response[1] = transactID[1];
		transactID[1] += 0x01;
		return response;
	}

	public byte[] getAddr() {
		return addr;
	}

	public byte[] getName() {
		return name;
	}
	
}