/*
 * @(#)SecureSocket.java
 *
 * Copyright (c) 2005-2006 JScape
 * 1147 S. 53rd Pl., Mesa, Arizona, 85206, U.S.A.
 * All rights reserved.
 *
 * This software is the confidential and proprietary information of
 * JScape. ("Confidential Information").  You shall not disclose such
 * Confidential Information and shall use it only in accordance with
 * the terms of the license agreement you entered into with JScape.
 */
package com.jscape.inet.ipclientssl;

import javax.net.ssl.*;
import java.io.InputStream;
import java.net.Socket;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.KeyStore;
import java.security.SecureRandom;
import java.security.cert.X509Certificate;

/**
 * Utility class for obtaining secure SSL/TLS sockets
 */
public class SecureSocket {

    protected static final String DEFAULT_PROTOCOL = "TLS";
    protected static final String DEFAULT_ALGORITHM = "SunX509";
    protected static final String DEFAULT_TYPE = "JKS";

    protected String protocol = DEFAULT_PROTOCOL;
    protected String algorithm = DEFAULT_ALGORITHM;
    protected String clientStoreType = DEFAULT_TYPE;
    protected String serverStoreType = DEFAULT_TYPE;

    protected KeyManager[] keyManagers = null;
    protected TrustManager[] trustManagers = getTrustAllManagers();
    private Socket socket = null;
    private SSLContext context = null;

    /**
     * Prevents this class from being instantiated
     *
     * @param socket  the socket
     * @param context the context
     */
    public SecureSocket(Socket socket, SSLContext context) {
        this.socket = socket;
        this.context = context;
    }

    /**
     * Secures a socket using the provided SSL context and socket.
     *
     * @return a secured Socket
     *
     * @throws Exception if an error occurs
     */
    public Socket getSocket()
            throws Exception {
        if (context == null) {
            context = getContext();
        }
        SSLSocketFactory factory = context.getSocketFactory();

        socket =
                factory.createSocket(
                        socket,
                        socket.getInetAddress().getHostAddress(),
                        socket.getPort(),
                        true);

        ((SSLSocket) socket).startHandshake();
        return socket;
    }

    /**
     * Gets trust all managers
     *
     * @return a TrustManager[]
     */
    protected static TrustManager[] getTrustAllManagers() {
        return new TrustManager[]{
                new X509TrustManager() {

                    /** @see X509TrustManager#getAcceptedIssuers() */
                    public X509Certificate[] getAcceptedIssuers() {
                        return new X509Certificate[0];
                    }

                    /** @see X509TrustManager#isClientTrusted(java.security.cert.X509Certificate[]) */
                    public boolean isClientTrusted(X509Certificate[] certs) {
                        return true;
                    }

                    /** @see X509TrustManager#isServerTrusted(java.security.cert.X509Certificate[]) */
                    public boolean isServerTrusted(X509Certificate[] certs) {
                        return true;
                    }

                    public void checkClientTrusted(
                            X509Certificate[] certs,
                            String val) {

                    }

                    public void checkServerTrusted(
                            X509Certificate[] certs,
                            String val) {

                    }
                }
        };
    }

    /**
     * Sets the client certificates
     *
     * @param filename        the keystore filename
     * @param password        the file password
     * @param clientStoreType the store type
     *
     * @throws Exception if an error occurs
     */
    public void setClientCertificates(String filename, String password, String clientStoreType)
            throws Exception {
        this.clientStoreType = clientStoreType != null ? clientStoreType : DEFAULT_TYPE;
        KeyStore keyStore = KeyStore.getInstance(this.clientStoreType);

        try (InputStream in = Files.newInputStream(Paths.get(filename))) {
            keyStore.load(in, password.toCharArray());
            KeyManagerFactory factory = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
            factory.init(keyStore, password.toCharArray());
            this.keyManagers = factory.getKeyManagers();
        }
    }

    /**
     * Sets client certificates
     *
     * @param filename the keystore filename
     * @param password the keystore password
     *
     * @throws Exception if an exception occurs
     */
    public void setClientCertificates(String filename, String password)
            throws Exception {
        setClientCertificates(filename, password, DEFAULT_TYPE);
    }

    /**
     * Sets the trusted server certificates.
     *
     * @param filename        the trust store filename
     * @param password        the store password
     * @param serverStoreType the store type
     *
     * @throws Exception if an error occurs
     */
    public void setServerCertificates(String filename, String password, String serverStoreType)
            throws Exception {
        this.serverStoreType = serverStoreType != null ? serverStoreType : DEFAULT_TYPE;
        KeyStore trustStore = KeyStore.getInstance(this.serverStoreType);

        try (InputStream in = Files.newInputStream(Paths.get(filename))) {
            trustStore.load(in, password.toCharArray());
            TrustManagerFactory factory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
            factory.init(trustStore);
            this.trustManagers = factory.getTrustManagers();
        }
    }

    /**
     * Sets server certificates
     *
     * @param filename the keystore filename
     * @param password the keystore password
     *
     * @throws Exception if an exception occurs
     */
    public void setServerCertificates(String filename, String password)
            throws Exception {
        setServerCertificates(filename, password, DEFAULT_TYPE);
    }

    /**
     * Gets an SSLContext
     *
     * @return an SSLContext
     *
     * @throws Exception if an error occurs
     */
    public SSLContext getContext()
            throws Exception {
        SSLContext context = SSLContext.getInstance(this.protocol);
        context.init(this.keyManagers, this.trustManagers, new SecureRandom());
        return context;
    }
}
