package ru.bloof.stun;

import java.io.IOException;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetSocketAddress;
import java.net.SocketTimeoutException;
import java.util.Arrays;
import java.util.Random;

import org.ice4j.StunException;
import org.ice4j.TransportAddress;
import org.ice4j.attribute.Attribute;
import org.ice4j.attribute.AttributeFactory;
import org.ice4j.attribute.ChangeRequestAttribute;
import org.ice4j.attribute.MappedAddressAttribute;
import org.ice4j.attribute.XorMappedAddressAttribute;
import org.ice4j.message.Message;
import org.ice4j.message.MessageFactory;
import org.ice4j.message.Request;

public class NatTypeTester {
    private static final byte[] buf = new byte[65536];
    private static final DatagramPacket recvPacket = new DatagramPacket(buf, buf.length);
    private static final Random rnd = new Random();
    private static final InetSocketAddress TEST = new InetSocketAddress("videostuntest.mycdn.me", 3478);

    public static void main(String[] args) throws Exception {
        try (DatagramSocket ds = new DatagramSocket(new InetSocketAddress(args[0], rnd.nextInt(10000) + 40000))) {
            ds.setSoTimeout(5000);
            Request test1 = createRequest();
            byte[] encoded = test1.encode(null);
            ds.send(new DatagramPacket(encoded, encoded.length, TEST));
            InetSocketAddress localAddress = (InetSocketAddress) ds.getLocalSocketAddress();
            System.out.println(localAddress);
            Message message;
            try {
                message = receive(ds, recvPacket, test1);
            } catch (SocketTimeoutException e) {
                System.out.println("No UDP for you");
                return;
            }
            TransportAddress ta = extractAddress(message);
            if (ta == null) {
                System.err.println("Invalid reponse");
                return;
            } else {
                System.out.println("Test1 mapped " + ta);
            }
            InetSocketAddress test1Address = (InetSocketAddress) recvPacket.getSocketAddress();
            if (localAddress.getPort() == ta.getPort() && localAddress.getAddress().equals(ta.getAddress())) {
                Request test2 = createRequest();
                ChangeRequestAttribute changeAttribute = AttributeFactory.createChangeRequestAttribute(true, true);
                test2.putAttribute(changeAttribute);
                byte[] test2Encoded = test2.encode(null);
                ds.send(new DatagramPacket(test2Encoded, test2Encoded.length, TEST));
                try {
                    Message receive = receive(ds, recvPacket, test2);
                    System.out.println("Test2 mapped " + extractAddress(receive));
                    InetSocketAddress test2Address = (InetSocketAddress) recvPacket.getSocketAddress();
                    if (test1Address.getAddress().equals(test2Address.getAddress())
                            || test1Address.getPort() == test2Address.getPort()) {
                        System.out.println("Server fooled us! Src=" + test1Address + " test2=" + test2Address);
                    }
                    System.out.println("Open internet");
                } catch (SocketTimeoutException e) {
                    System.out.println("Symmetric NAT, also changing MAPPED-ADDRESS");
                }
            } else {
                Request test2 = createRequest();
                ChangeRequestAttribute changeAttribute = AttributeFactory.createChangeRequestAttribute(true, true);
                test2.putAttribute(changeAttribute);
                byte[] test2Encoded = test2.encode(null);
                ds.send(new DatagramPacket(test2Encoded, test2Encoded.length, TEST));
                try {
                    Message receive = receive(ds, recvPacket, test2);
                    System.out.println("Test2 mapped " + extractAddress(receive));
                    InetSocketAddress test2Address = (InetSocketAddress) recvPacket.getSocketAddress();
                    if (test1Address.getAddress().equals(test2Address.getAddress())
                            || test1Address.getPort() == test2Address.getPort()) {
                        System.out.println("Server fooled us! Src=" + test1Address + " test2=" + test2Address);
                    }
                    System.out.println("NAT, no filter");
                    return;
                } catch (SocketTimeoutException e) {
                }

                Request test3 = createRequest();
                ChangeRequestAttribute portChangeAttribute = AttributeFactory.createChangeRequestAttribute(false,
                        true);
                test3.putAttribute(portChangeAttribute);
                byte[] test3Encode = test3.encode(null);
                ds.send(new DatagramPacket(test3Encode, test3Encode.length, TEST));
                try {
                    Message receive = receive(ds, recvPacket, test3);
                    System.out.println("Test3 mapped " + extractAddress(receive));
                    InetSocketAddress test3Address = (InetSocketAddress) recvPacket.getSocketAddress();
                    if (test1Address.getPort() == test3Address.getPort()) {
                        System.out.println("Server fooled us! Src=" + test1Address + " test3=" + test3Address);
                    }
                    System.out.println("NAT, address-sensitive filter");
                } catch (SocketTimeoutException e) {
                    System.out.println("NAT, port-sensitive filter");
                }
            }
        }
    }

    private static Request createRequest() throws StunException {
        Request request = MessageFactory.createBindingRequest();
        byte[] tId = new byte[12];
        rnd.nextBytes(tId);
        request.setTransactionID(tId);
        return request;
    }

    private static TransportAddress extractAddress(Message message) {
        Attribute attr = message.getAttribute(Attribute.MAPPED_ADDRESS);
        if (attr != null) {
            return ((MappedAddressAttribute) attr).getAddress();
        }
        attr = message.getAttribute(Attribute.XOR_MAPPED_ADDRESS);
        if (attr != null) {
            return ((XorMappedAddressAttribute) attr).getAddress(message.getTransactionID());
        }
        return null;
    }

    private static Message receive(DatagramSocket ds, DatagramPacket packet, Request request)
            throws IOException, StunException {
        while (true) {
            ds.receive(packet);
            Message decoded = Message.decode(packet.getData(), (char) packet.getOffset(), (char) packet.getLength());
            if (Arrays.equals(decoded.getTransactionID(), request.getTransactionID())) {
                System.out.println(
                        "Received success " + (decoded.getMessageType() == Message.BINDING_SUCCESS_RESPONSE));
                return decoded;
            }
        }
    }
}