// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.
package streamer;

import static rdpclient.MockServer.Packet.PacketType.CLIENT;
import static rdpclient.MockServer.Packet.PacketType.SERVER;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.util.HashMap;

import javax.net.SocketFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLSocket;
import javax.net.ssl.SSLSocketFactory;
import javax.net.ssl.TrustManager;

import rdpclient.MockServer;
import rdpclient.MockServer.Packet;
import rdpclient.TrustAllX509TrustManager;

public class SocketWrapper extends PipelineImpl {

    protected InputStreamSource source;
    protected OutputStreamSink sink;
    protected Socket socket;
    protected InetSocketAddress address;

    protected SSLSocket sslSocket;

    //protected String SSL_VERSION_TO_USE = "TLSv1.2";
    /*DEBUG*/protected String SSL_VERSION_TO_USE = "TLSv1";

    public SocketWrapper(String id) {
        super(id);
    }

    @Override
    protected HashMap<String, Element> initElementMap(String id) {
        HashMap<String, Element> map = new HashMap<String, Element>();

        source = new InputStreamSource(id + "." + OUT, this);
        sink = new OutputStreamSink(id + "." + IN, this);

        // Pass requests to read data to socket input stream
        map.put(OUT, source);

        // All incoming data, which is sent to this socket wrapper, will be sent
        // to socket remote
        map.put(IN, sink);

        return map;
    }

    /**
     * Connect this socket wrapper to remote server and start main loop on
     * IputStreamSource stdout link, to watch for incoming data, and
     * OutputStreamSink stdin link, to pull for outgoing data.
     *
     * @param address
     * @throws IOException
     */
    public void connect(InetSocketAddress address) throws IOException {
        this.address = address;

        // Connect socket to server
        socket = SocketFactory.getDefault().createSocket();
        try {
            socket.connect(address);

            InputStream is = socket.getInputStream();
            source.setInputStream(is);

            OutputStream os = socket.getOutputStream();
            sink.setOutputStream(os);

            // Start polling for data to send to remote sever
            runMainLoop(IN, STDIN, true, true);

            // Push incoming data from server to handlers
            runMainLoop(OUT, STDOUT, false, false);

        } finally {
            socket.close();
        }
    }

    @Override
    public void handleEvent(Event event, Direction direction) {
        switch (event) {
            case SOCKET_UPGRADE_TO_SSL:
                upgradeToSsl();
                break;
            default:
                super.handleEvent(event, direction);
                break;
        }
    }

    public void upgradeToSsl() {

        if (sslSocket != null)
            // Already upgraded
            return;

        if (verbose)
            System.out.println("[" + this + "] INFO: Upgrading socket to SSL.");

        try {
            // Use most secure implementation of SSL available now.
            // JVM will try to negotiate TLS1.2, then will fallback to TLS1.0, if
            // TLS1.2 is not supported.
            SSLContext sslContext = SSLContext.getInstance(SSL_VERSION_TO_USE);

            // Trust all certificates (FIXME: insecure)
            sslContext.init(null, new TrustManager[] {new TrustAllX509TrustManager()}, null);

            SSLSocketFactory sslSocketFactory = sslContext.getSocketFactory();
            sslSocket = (SSLSocket)sslSocketFactory.createSocket(socket, address.getHostName(), address.getPort(), true);
            sslSocket.startHandshake();

            InputStream sis = sslSocket.getInputStream();
            source.setInputStream(sis);

            OutputStream sos = sslSocket.getOutputStream();
            sink.setOutputStream(sos);

        } catch (Exception e) {
            throw new RuntimeException("Cannot upgrade socket to SSL: " + e.getMessage(), e);
        }

    }

    @Override
    public void validate() {
        for (Element element : elements.values())
            element.validate();

        if (get(IN).getPads(Direction.IN).size() == 0)
            throw new RuntimeException("[ " + this + "] Input of socket is not connected.");

        if (get(OUT).getPads(Direction.OUT).size() == 0)
            throw new RuntimeException("[ " + this + "] Output of socket is not connected.");

    }

    public void shutdown() {
        try {
            handleEvent(Event.STREAM_CLOSE, Direction.IN);
        } catch (Exception e) {
        }
        try {
            handleEvent(Event.STREAM_CLOSE, Direction.OUT);
        } catch (Exception e) {
        }
        try {
            if (sslSocket != null)
                sslSocket.close();
        } catch (Exception e) {
        }
        try {
            socket.close();
        } catch (Exception e) {
        }
    }

    @Override
    public String toString() {
        return "SocketWrapper(" + id + ")";
    }

    /**
     * Example.
     */
    public static void main(String args[]) {
        try {
            System.setProperty("streamer.Link.debug", "true");
            System.setProperty("streamer.Element.debug", "true");
            System.setProperty("rdpclient.MockServer.debug", "true");

            Pipeline pipeline = new PipelineImpl("echo client");

            SocketWrapper socketWrapper = new SocketWrapper("socket");

            pipeline.add(socketWrapper);
            pipeline.add(new BaseElement("echo"));

            pipeline.link("socket", "echo", "socket");

            final byte[] mockData = new byte[] {0x01, 0x02, 0x03};
            MockServer server = new MockServer(new Packet[] {new Packet("Server hello") {
                {
                    type = SERVER;
                    data = mockData;
                }
            }, new Packet("Client hello") {
                {
                    type = CLIENT;
                    data = mockData;
                }
            }, new Packet("Server hello") {
                {
                    type = SERVER;
                    data = mockData;
                }
            }, new Packet("Client hello") {
                {
                    type = CLIENT;
                    data = mockData;
                }
            }});
            server.start();
            InetSocketAddress address = server.getAddress();

            socketWrapper.connect(address);

        } catch (IOException e) {
            e.printStackTrace(System.err);
        }

    }
}