//
//  This program implements a simple DNS relay server.
//
//  This program is for demonstration purposes within the Computer
//  Networking in the Faculty of informatics of Università della
//  Svizzera italiana.  See http://www.inf.usi.ch/carzaniga/edu/ntw/
//
//  Author: Antonio Carzaniga (antonio.carzaniga@usi.ch)
//
//  Copyright (C) 2018 Università della Svizzera italiana
//
//  This is free software: you can redistribute it and/or modify
//  it under the terms of the GNU General Public License as published by
//  the Free Software Foundation, either version 3 of the License, or
//  (at your option) any later version.
//
//  Siena is distributed in the hope that it will be useful,
//  but WITHOUT ANY WARRANTY; without even the implied warranty of
//  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//  GNU General Public License for more details.
//
//  See <http://www.gnu.org/licenses/>.
//
import java.net.DatagramSocket;
import java.net.InetAddress;
import java.net.SocketAddress;
import java.net.InetSocketAddress;
import java.net.DatagramPacket;
import java.util.HashMap;

/** Utility functions to read data from a packet.
 *
 *  A packet is just a sequence of "octects", meaning 8-bit data
 *  units, meaning bytes in Java.  As a general principle, we read
 *  data starting at the given offset from the given buffer.  The
 *  numeric data is read in "network order", which is equivalent to
 *  big-endian format.  (See
 *  https://en.wikipedia.org/wiki/Endianness#Big)
 */
class InetFormat {
    static int read_int32(byte [] buffer, int offset) {
	int result = read_byte(buffer, offset++);
	result = (result << 8) | read_byte(buffer, offset++);
	result = (result << 8) | read_byte(buffer, offset++);
	result = (result << 8) | read_byte(buffer, offset);
	return result;
    }

    static long read_uint32(byte [] buffer, int offset) {
	long result = read_byte(buffer, offset++);
	result = (result << 8) | read_byte(buffer, offset++);
	result = (result << 8) | read_byte(buffer, offset++);
	result = (result << 8) | read_byte(buffer, offset);
	return result;
    }

    static int read_uint16(byte [] buffer, int offset) {
	return (read_byte(buffer, offset) << 8) | read_byte(buffer, offset + 1);
    }

    static boolean read_bool(byte [] buffer, int byte_offset, int bit_offset) {
	byte mask = 1;
	mask <<= (7 - bit_offset);
	return (buffer[byte_offset] & mask) != 0;
    }

    static byte read_bitfield(byte [] buffer, int byte_offset, int bit_offset, int length) {
	byte mask = ~0;
	mask >>= (8 - length);
	mask &= buffer[byte_offset] >> (8 - length - bit_offset);
	return mask;
    }

    static int read_byte(byte [] buffer, int offset) {
	return buffer[offset] & 0xff;
    }
}

/** A Name as a sequence of labels, as defined in Section 3.1 of RFC 1035
 */
class Name extends InetFormat {
    byte [][] labels;

    /** read a Name starting at the offset within the given buffer.
     * 
     * @return the offset of the element in the buffer immediately following the Name
     */
    public int read(byte [] buffer, int offset) {
	// a potential "pointer" to a previous name (see Section 4.1.4)
	int pointer_offset = -1;
	if (read_byte(buffer,offset) >= 192) {
	    pointer_offset = offset + 2;
	    offset = read_uint16(buffer, offset) & 0x3fff;
	}
	int count = 0;
	int i = offset;
	int llen = read_byte(buffer, i);
	while (llen > 0) {
	    i += llen + 1;
	    llen = read_byte(buffer, i);
	    count += 1;
	}
	labels = new byte[count][];
	for (i = 0; i < count; ++i) {
	    llen = read_byte(buffer, offset);
	    offset += 1;
	    labels[i] = java.util.Arrays.copyOfRange(buffer, offset, offset + llen);
	    offset += llen;
	}
	offset += 1;
	return (pointer_offset < 0) ? offset : pointer_offset;
    }

    public String toString() {
	int len = labels.length;
	for (int i = 0; i < labels.length; ++i) 
	    len += labels[i].length;
	byte [] name = new byte[len];
	int k = 0;
	for (int i = 0; i < labels.length; ++i) {
	    for (int j = 0; j < labels[i].length; ++j) 
		name[k++] = labels[i][j];
	    name[k++] = (byte)0x2e; /* '.' */
	}
	return new String(name);
    }
}

/** A Question in a DNS packet, as specified in Section 4.1.2 of RFC 1035.
 */
class Question extends InetFormat {
    Name name;
    int type;
    int clss;

    /** read a Question starting at the given offset within the given buffer.
     * 
     * @return the offset of the element in the buffer immediately following the Question
     */
    public int read(byte [] buffer, int offset) {
	name = new Name();
	offset = name.read(buffer, offset);
	type = read_uint16(buffer, offset);
	offset += 2;
	clss = read_uint16(buffer, offset);
	offset += 2;
	return offset;
    }

    public void print() {
	System.out.println("Question:"
			   + " name=" + name.toString()
			   + " type=" + type
			   + " clss=" + clss);
    }
};

/** A Resource Record in a DNS packet, as specified in Section 4.1.3 of RFC 1035.
 */
class RR extends InetFormat {
    Name name;
    int type;
    int clss;
    long ttl;
    int rdlength;
    byte [] rdata;

    /** read a Resource Record starting at the given offset within the given buffer.
     * 
     * @return the offset of the element in the buffer immediately following the RR
     */
    public int read(byte [] buffer, int offset) {
	name = new Name();
	offset = name.read(buffer, offset);
	type = read_uint16(buffer, offset);
	offset += 2;
	clss = read_uint16(buffer, offset);
	offset += 2;
	ttl = read_uint32(buffer, offset);
	offset += 2;
	rdlength = read_uint16(buffer, offset);
	offset += 2;
	rdata = java.util.Arrays.copyOfRange(buffer, offset, offset + rdlength);
	offset += rdlength;
	return offset;
    }

    public void print() {
	System.out.println("RR:"
			   + " name=" + name.toString()
			   + " type=" + type
			   + " clss=" + clss
			   + " ttl=" + ttl
			   + " rdlength=" + rdlength);
    }
};

/** DNS Message, as specified in Section 4.1 of RFC 1035.
 */
class DNSMessage extends InetFormat {
    int id;
    boolean qr;
    byte opcode;
    boolean aa;
    boolean tc;
    boolean rd;
    boolean ra;
    byte z;
    byte rcode;
    
    int qdcount;
    int ancount;
    int nscount;
    int arcount;

    Question [] qd;
    RR [] an;
    RR [] ns;

    public void print() {
	System.out.println("DNS packet:"
			   + " id=" + id 
			   + " qr=" + qr 
			   + " opcode=" + opcode 
			   + " aa=" + aa 
			   + " tc=" + tc 
			   + " rd=" + rd 
			   + " ra=" + ra 
			   + " z=" + z 
			   + " rcode=" + rcode 
			   + " qdcount=" + qdcount 
			   + " ancount=" + ancount 
			   + " nscount=" + nscount 
			   + " arcount=" + arcount);
	if (qdcount > 0) {
	    for (int i = 0; i < qdcount; ++i)
		qd[i].print();
	}
	if (ancount > 0) {
	    for (int i = 0; i < ancount; ++i)
		an[i].print();
	}
	if (nscount > 0) {
	    for (int i = 0; i < nscount; ++i)
		ns[i].print();
	}
    }
    
    /** read a DNS Message starting at the given offset within the given buffer.
     */
    public void read(byte [] buffer, int offset) {
	id = read_uint16(buffer, offset);
	offset += 2;
	qr = read_bool(buffer, offset, 0);
	opcode = read_bitfield(buffer, offset + 2, 1, 4);
	aa = read_bool(buffer, offset, 5);
	tc = read_bool(buffer, offset, 6);
	rd = read_bool(buffer, offset, 7);
	offset += 1;
	ra = read_bool(buffer, offset, 0);
	z = read_bitfield(buffer, offset, 1, 3);
	rcode = read_bitfield(buffer, offset, 4, 4);
	offset += 1;
	qdcount = read_uint16(buffer, offset);
	offset += 2;
	ancount = read_uint16(buffer, offset);
	offset += 2;
	nscount = read_uint16(buffer, offset);
	offset += 2;
	arcount = read_uint16(buffer, offset);
	offset += 2;
	if (qdcount > 0) {
	    qd = new Question[qdcount];
	    for (int i = 0; i < qdcount; ++i) {
		qd[i] = new Question();
		offset = qd[i].read(buffer, offset);
	    }
	}
	if (ancount > 0) {
	    an = new RR[ancount];
	    for (int i = 0; i < ancount; ++i) {
		an[i] = new RR();
		offset = an[i].read(buffer, offset);
	    }
	}
	if (nscount > 0) {
	    ns = new RR[nscount];
	    for (int i = 0; i < nscount; ++i) {
		ns[i] = new RR();
		offset = ns[i].read(buffer, offset);
	    }
	}
    }
}

/** A simple DNS "relay" server.
 *
 *  Serves DNS queries by simply relaying them to a root server.
 */
public class LocalDNS {
    static public void main(String [] args) {
	if (args.length != 2) {
	    System.err.println("usage: LocalDNS <port-number> <root-server-address>");
	    System.exit(1);
	}
	int local_port;
	SocketAddress root_address;
	try {
	    local_port = Integer.parseInt(args[0]);
	    root_address = new InetSocketAddress(args[1], 53);
	    try {
		byte [] buffer = new byte[65536];
		DatagramPacket pkt = new DatagramPacket(buffer, buffer.length);
		DatagramSocket socket = new DatagramSocket(local_port);

		HashMap<Integer, SocketAddress> pending = new HashMap<Integer, SocketAddress> ();
		DNSMessage message = new DNSMessage();

		while (true) {
		    socket.receive(pkt);
		    message.read(pkt.getData(), pkt.getOffset());
		    message.print();
		    if (message.qr == false /* QUERY */) {
			SocketAddress client_addr = pkt.getSocketAddress();
			System.out.println("query from " + client_addr + " --> root server");
			pending.put(message.id, client_addr);
			pkt.setSocketAddress(root_address);
			socket.send(pkt);
		    } else /* RESPONSE */ {
			System.out.print("reply from " + pkt.getSocketAddress());
			SocketAddress client_addr = pending.remove(message.id);
			if (client_addr == null) {
			    System.out.println(": id " + message.id + " not found!");
			} else {
			    System.out.println(" --> " + client_addr);
			    pkt.setSocketAddress(client_addr);
			    socket.send(pkt);
			}
		    }
		}
	    } catch (Exception ex) {
		System.err.println(ex);
		System.err.println("Terminating LocalDNS loop.");
	    }
	} catch (Exception ex) {
	    System.err.println(ex);
	    System.err.println("usage: LocalDNS <port-number> <root-server-address>");
	    System.exit(2);
	}
    }
}