package net.sourceforge.jsocks;
import java.net.*;
import java.io.*;

/**
  Datagram socket to interract through the firewall.<BR>
  Can be used same way as the normal DatagramSocket. One should
  be carefull though with the datagram sizes used, as additional data
  is present in both incomming and outgoing datagrams.
   <p>
   SOCKS5 protocol allows to send host address as either:
   <ul>
    <li> IPV4, normal 4 byte address. (10 bytes header size)
    <li> IPV6, version 6 ip address (not supported by Java as for now).
         22 bytes header size.
    <li> Host name,(7+length of the host name bytes header size).
   </ul>
  As with other Socks equivalents, direct addresses are handled
  transparently, that is data will be send directly when required
  by the proxy settings.
  <p>
  <b>NOTE:</b><br>
  Unlike other SOCKS Sockets, it <b>does not</b> support proxy chaining,
  and will throw an exception if proxy has a chain proxy attached. The
  reason for that is not my laziness, but rather the restrictions of
  the SOCKSv5 protocol. Basicaly SOCKSv5 proxy server, needs to know from
  which host:port datagrams will be send for association, and returns address
  to which datagrams should be send by the client, but it does not
  inform client from which host:port it is going to send datagrams, in fact
  there is even no guarantee they will be send at all and from the same address
  each time.

 */
public class Socks5DatagramSocket extends DatagramSocket {

    InetAddress relayIP;
    int relayPort;
    Socks5Proxy proxy;
    private boolean server_mode = false;
    UDPEncapsulation encapsulation;


    /**
       Construct Datagram socket for communication over SOCKS5 proxy
       server. This constructor uses default proxy, the one set with
       Proxy.setDefaultProxy() method. If default proxy is not set or
       it is set to version4 proxy, which does not support datagram
       forwarding, throws SocksException.

     */
    public Socks5DatagramSocket() throws SocksException,
        IOException {
        this(Proxy.defaultProxy, 0, null);
    }
    /**
       Construct Datagram socket for communication over SOCKS5 proxy
       server. And binds it to the specified local port.
       This constructor uses default proxy, the one set with
       Proxy.setDefaultProxy() method. If default proxy is not set or
       it is set to version4 proxy, which does not support datagram
       forwarding, throws SocksException.
     */
    public Socks5DatagramSocket(int port) throws SocksException,
        IOException {
        this(Proxy.defaultProxy, port, null);
    }
    /**
       Construct Datagram socket for communication over SOCKS5 proxy
       server. And binds it to the specified local port and address.
       This constructor uses default proxy, the one set with
       Proxy.setDefaultProxy() method. If default proxy is not set or
       it is set to version4 proxy, which does not support datagram
       forwarding, throws SocksException.
     */
    public Socks5DatagramSocket(int port, InetAddress ip) throws SocksException,
        IOException {
        this(Proxy.defaultProxy, port, ip);
    }

    /**
      Constructs datagram socket for communication over specified proxy.
      And binds it to the given local address and port. Address of null
      and port of 0, signify any availabale port/address.
      Might throw SocksException, if:
      <ol>
       <li> Given version of proxy does not support UDP_ASSOCIATE.
       <li> Proxy can't be reached.
       <li> Authorization fails.
       <li> Proxy does not want to perform udp forwarding, for any reason.
      </ol>
      Might throw IOException if binding dtagram socket to given address/port
      fails.
      See java.net.DatagramSocket for more details.
     */
    public Socks5DatagramSocket(Proxy p, int port, InetAddress ip)
    throws SocksException,
        IOException {
        super(port, ip);

        if (p == null) throw new SocksException(Proxy.SOCKS_NO_PROXY);

        if (!(p instanceof Socks5Proxy))
            throw new SocksException(-1, "Datagram Socket needs Proxy version 5");

        if (p.chainProxy != null)
            throw new SocksException(Proxy.SOCKS_JUST_ERROR,
                                     "Datagram Sockets do not support proxy chaining.");

        proxy = (Socks5Proxy) p.copy();
        ProxyMessage msg = proxy.udpAssociate(super.getLocalAddress(),
                                              super.getLocalPort());
        relayIP = msg.ip;

        if (relayIP.getHostAddress().equals("0.0.0.0")) relayIP = proxy.proxyIP;

        relayPort = msg.port;
        encapsulation = proxy.udp_encapsulation;
        //debug("Datagram Socket:"+getLocalAddress()+":"+getLocalPort()+"\n");
        //debug("Socks5Datagram: "+relayIP+":"+relayPort+"\n");
    }

    /**
      Used by UDPRelayServer.
     */
    Socks5DatagramSocket(boolean server_mode, UDPEncapsulation encapsulation,
                         InetAddress relayIP, int relayPort)
    throws IOException {
        super();
        this.server_mode = server_mode;
        this.relayIP = relayIP;
        this.relayPort = relayPort;
        this.encapsulation = encapsulation;
        this.proxy = null;
    }

    /**
      Sends the Datagram either through the proxy or directly depending
      on current proxy settings and destination address. <BR>

      <B> NOTE: </B> DatagramPacket size should be at least 10 bytes less
                     than the systems limit.

      <P>
      See documentation on java.net.DatagramSocket
      for full details on how to use this method.
      @param dp Datagram to send.
      @throws IOException If error happens with I/O.
     */
    public void send(DatagramPacket dp) throws IOException {
        //If the host should be accessed directly, send it as is.
        if (!server_mode) {
            super.send(dp);
            //debug("Sending directly:");
            return;
        }

        byte[] head = formHeader(dp.getAddress(), dp.getPort());
        byte[] buf = new byte[head.length + dp.getLength()];
        byte[] data = dp.getData();
        //Merge head and data
        System.arraycopy(head, 0, buf, 0, head.length);
        //System.arraycopy(data,dp.getOffset(),buf,head.length,dp.getLength());
        System.arraycopy(data, 0, buf, head.length, dp.getLength());

        if (encapsulation != null)
            buf = encapsulation.udpEncapsulate(buf, true);

        super.send(new DatagramPacket(buf, buf.length, relayIP, relayPort));
    }
    /**
      This method allows to send datagram packets with address type DOMAINNAME.
      SOCKS5 allows to specify host as names rather than ip addresses.Using
      this method one can send udp datagrams through the proxy, without having
      to know the ip address of the destination host.
      <p>
      If proxy specified for that socket has an option resolveAddrLocally set
      to true host will be resolved, and the datagram will be send with address
      type IPV4, if resolve fails, UnknownHostException is thrown.
      @param dp Datagram to send, it should contain valid port and data
      @param host Host name to which datagram should be send.
      @throws IOException If error happens with I/O, or the host can't be
      resolved when proxy settings say that hosts should be resolved locally.
      @see Socks5Proxy#resolveAddrLocally(boolean)
     */
    public void send(DatagramPacket dp, String host) throws IOException {
        dp.setAddress(InetAddress.getByName(host));
        super.send(dp);
    }

    /**
     * Receives udp packet. If packet have arrived from the proxy relay server,
     * it is processed and address and port of the packet are set to the
     * address and port of sending host.<BR>
     * If the packet arrived from anywhere else it is not changed.<br>
     * <B> NOTE: </B> DatagramPacket size should be at least 10 bytes bigger
     * than the largest packet you expect (this is for IPV4 addresses).
     * For hostnames and IPV6 it is even more.
       @param dp Datagram in which all relevent information will be copied.
     */
    public void receive(DatagramPacket dp) throws IOException {
        super.receive(dp);

        if (server_mode) {
            //Drop all datagrams not from relayIP/relayPort
            int init_length = dp.getLength();
            int initTimeout = getSoTimeout();
            long startTime = System.currentTimeMillis();

            while (!relayIP.equals(dp.getAddress()) ||
                    relayPort != dp.getPort()) {
                //Restore datagram size
                dp.setLength(init_length);

                //If there is a non-infinit timeout on this socket
                //Make sure that it happens no matter how often unexpected
                //packets arrive.
                if (initTimeout != 0) {
                    int newTimeout = initTimeout - (int)(System.currentTimeMillis() -
                                                         startTime);

                    if (newTimeout <= 0) throw new InterruptedIOException(
                            "In Socks5DatagramSocket->receive()");

                    setSoTimeout(newTimeout);
                }

                super.receive(dp);
            }

            //Restore timeout settings
            if (initTimeout != 0) setSoTimeout(initTimeout);
        }
        else if (!relayIP.equals(dp.getAddress()) ||
                 relayPort != dp.getPort())
            return; // Recieved direct packet

        //If the datagram is not from the relay server, return it it as is.
        byte[] data;
        data = dp.getData();

        if (encapsulation != null)
            data = encapsulation.udpEncapsulate(data, false);

        int offset = 0; //Java 1.1
        //int offset = dp.getOffset(); //Java 1.2
        ByteArrayInputStream bIn = new ByteArrayInputStream(data, offset,
                dp.getLength());
        ProxyMessage msg = new Socks5Message(bIn);
        dp.setPort(msg.port);
        dp.setAddress(msg.getInetAddress());
        //what wasn't read by the Message is the data
        int data_length = bIn.available();
        //Shift data to the left
        System.arraycopy(data, offset + dp.getLength() - data_length,
                         data, offset, data_length);
        dp.setLength(data_length);
    }

    /**
     * Returns port assigned by the proxy, to which datagrams are relayed.
     * It is not the same port to which other party should send datagrams.
       @return Port assigned by socks server to which datagrams are send
       for association.
     */
    public int getLocalPort() {
        if (server_mode) return super.getLocalPort();

        return relayPort;
    }
    /**
     * Address assigned by the proxy, to which datagrams are send for relay.
     * It is not necesseraly the same address, to which other party should send
     * datagrams.
       @return Address to which datagrams are send for association.
     */
    public InetAddress getLocalAddress() {
        if (server_mode) return super.getLocalAddress();

        return relayIP;
    }

    /**
     * Closes datagram socket, and proxy connection.
     */
    public void close() {
        if (!server_mode) proxy.endSession();

        super.close();
    }

    /**
      This method checks wether proxy still runs udp forwarding service
      for this socket.
      <p>
      This methods checks wether the primary connection to proxy server
      is active. If it is, chances are that proxy continues to forward
      datagrams being send from this socket. If it was closed, most likely
      datagrams are no longer being forwarded by the server.
      <p>
      Proxy might decide to stop forwarding datagrams, in which case it
      should close primary connection. This method allows to check, wether
      this have been done.
      <p>
      You can specify timeout for which we should be checking EOF condition
      on the primary connection. Timeout is in milliseconds. Specifying 0 as
      timeout implies infinity, in which case method will block, until
      connection to proxy is closed or an error happens, and then return false.
      <p>
      One possible scenario is to call isProxyactive(0) in separate thread,
      and once it returned notify other threads about this event.

      @param timeout For how long this method should block, before returning.
      @return true if connection to proxy is active, false if eof or error
              condition have been encountered on the connection.
    */
    public boolean isProxyAlive(int timeout) {
        if (server_mode) return false;

        if (proxy != null) {
            try {
                proxy.proxySocket.setSoTimeout(timeout);
                int eof = proxy.in.read();

                if (eof < 0) return false; // EOF encountered.
                else return true;         // This really should not happen
            }
            catch (InterruptedIOException iioe) {
                return true;          // read timed out.
            }
            catch (IOException ioe) {
                return false;
            }
        }

        return false;
    }

//PRIVATE METHODS
//////////////////


    private byte[] formHeader(InetAddress ip, int port) {
        Socks5Message request = new Socks5Message(0, ip, port);
        request.data[0] = 0;
        return request.data;
    }


    /*======================================================================

    //Mainly Test functions
    //////////////////////

       private String bytes2String(byte[] b){
          String s="";
          char[] hex_digit = { '0','1','2','3','4','5','6','7','8','9',
                               'A','B','C','D','E','F'};
          for(int i=0;i<b.length;++i){
              int i1 = (b[i] & 0xF0) >> 4;
              int i2 = b[i] & 0xF;
              s+=hex_digit[i1];
              s+=hex_digit[i2];
              s+=" ";
          }
          return s;
       }
       private static final void debug(String s){
          if(DEBUG)
             System.out.print(s);
       }

       private static final boolean DEBUG = true;


       public static void usage(){
          System.err.print(
        "Usage: java Socks.SocksDatagramSocket host port [socksHost socksPort]\n");
       }

       static final int defaultProxyPort = 1080;           //Default Port
       static final String defaultProxyHost = "www-proxy"; //Default proxy

       public static void main(String args[]){
          int port;
          String host;
          int proxyPort;
          String proxyHost;
          InetAddress ip;

          if(args.length > 1 && args.length < 5){
         try{

             host = args[0];
             port = Integer.parseInt(args[1]);

             proxyPort =(args.length > 3)? Integer.parseInt(args[3])
                                         : defaultProxyPort;

             host = args[0];
             ip = InetAddress.getByName(host);

             proxyHost =(args.length > 2)? args[2]
                                         : defaultProxyHost;

             Proxy.setDefaultProxy(proxyHost,proxyPort);
             Proxy p = Proxy.getDefaultProxy();
             p.addDirect("lux");


             DatagramSocket ds = new Socks5DatagramSocket();


             BufferedReader in = new BufferedReader(
                     new InputStreamReader(System.in));
                 String s;

                 System.out.print("Enter line:");
                 s = in.readLine();

             while(s != null){
                    byte[] data = (s+"\r\n").getBytes();
                    DatagramPacket dp = new DatagramPacket(data,0,data.length,
                                            ip,port);
                    System.out.println("Sending to: "+ip+":"+port);
                    ds.send(dp);
                dp = new DatagramPacket(new byte[1024],1024);

                System.out.println("Trying to recieve on port:"+
                                    ds.getLocalPort());
                ds.receive(dp);
                System.out.print("Recieved:\n"+
                                 "From:"+dp.getAddress()+":"+dp.getPort()+
                                 "\n\n"+
                    new String(dp.getData(),dp.getOffset(),dp.getLength())+"\n"
                    );
                    System.out.print("Enter line:");
                    s = in.readLine();

             }
             ds.close();
             System.exit(1);

         }catch(SocksException s_ex){
           System.err.println("SocksException:"+s_ex);
           s_ex.printStackTrace();
           System.exit(1);
         }catch(IOException io_ex){
           io_ex.printStackTrace();
           System.exit(1);
         }catch(NumberFormatException num_ex){
           usage();
           num_ex.printStackTrace();
           System.exit(1);
         }

          }else{
        usage();
          }
       }
    */

}