package com.microfocus.ucmdb.passwordvault.spi;

import com.microfocus.ucmdb.passwordvault.spi.impl.CommonPasswordVaultParameters;
import org.json.JSONException;
import org.json.JSONObject;

import javax.net.ssl.*;
import java.io.*;
import java.net.HttpURLConnection;
import java.net.URL;
import java.security.KeyStore;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

/**
 * These three @Override methods are mandatory.
 */
public class HashiCorpPasswordVault extends CommonPasswordVaultParameters implements PasswordVault {
    private static String vaultFilePath = System.getProperty("java.home") + File.separator + ".." + File.separator + ".." + File.separator + "runtime" + File.separator + "probeManager" + File.separator + "passwordVaults/hashiCorp/";
    /* The SSL context */
    private static SSLContext sc;
    /* The URL of vault API server */
    private String HashiCorpAPIServer;


    /***********
     * Authentication mode. "1" is username and password,  "2" is root token and AppRole
     ******************/
    private String AuthenticationMode;

    /***********
     * username and password authentication.
     ******************/
    /* The username of vault server */
    private String HashiCorpUsername;
    /* The password of vault server */
    private String HashiCorpPassword;
    /* The password of truststore */
    private String TrustStorePassword;
    /* The root access token of vault */
    private static String clientToken;


    /***********
     * root token and appRole
     ******************/
    /* The root access token of vault */
    private String RootToken;
    /* The available appRole of vault server */
    private String AppRole;
    /* the cache of aws access info.  the data structure is :  key is referenceId, value is json map, it includes awsLeaseId,awsSecretKey and awsAccessKey */
    private JSONObject awsAllAccessCache;

    /**
     * This method can load parameter values from your confiuration file.
     *
     * @param parameters These parameters are from your configuration file.
     */
    @Override
    public void setParameters(Map<String, String> parameters) {
        super.setParameters(parameters);
        HashiCorpAPIServer = this.getParameter("HashiCorpAPIServer");
        TrustStorePassword = this.getParameter("TrustStorePassword");

        AuthenticationMode = this.getParameter("AuthenticationMode");

        HashiCorpUsername = this.getParameter("HashiCorpUsername");
        HashiCorpPassword = this.getParameter("HashiCorpPassword");

        RootToken = this.getParameter("RootToken");
        AppRole = this.getParameter("AppRole");

        if (this.getParameter("awsAllAccessCache") != null) {
            try {
                awsAllAccessCache = new JSONObject(this.getParameter("awsAllAccessCache"));
            } catch (JSONException e) {
            }
        } else {
            awsAllAccessCache = new JSONObject();
        }
    }

    /**
     * @return the provider name, and it must equal the value that you put in JMX passwordVault.provider.list
     */
    @Override
    public String getProviderName() {
        return "HashiCorp";
    }

    /**
     * this is the main method to get field values from external vault.
     *
     * @param referenceID       This is the reference ID in UCMDB credential and is mapped to the path in HashiCorp vault.
     * @param attributesInVault These are attributes mapped between vault and UCMDB credential, e.g. ["name":"protocol_username","pwd":"protocol_password"]
     * @param protocolEntry     This the credential from UCMDB.
     * @return true, if input data is valid; false, if input data in invalid.
     * @throws PasswordVaultException
     */
    @Override
    public boolean retrieveAndPopulateVaultAttributes(String referenceID, List<String> attributesInVault, Map<String, String> protocolEntry) throws PasswordVaultException {
        //Validate the input data.
        if (referenceID == null || referenceID.isEmpty()) {
            return false;
        }
        //If vault server is in SSL mode, initialize the trust store chain at first.
        if (sc == null) {
            initTrustStore();
        }

        //check if the client token is saved in cache or not.
        if (clientToken == null) {
            //If not, authenticate it to get the token stored here.
            if ("1".equals(AuthenticationMode)) {
                authenticateUserPwd();
            } else if ("2".equals(AuthenticationMode)) {
                authenticateByAppRole();
            } else {
                throw new PasswordVaultException("Authentication mode is not set, please set 1 for userpwd  or  2 for appRole.");
            }
        }

        Map<String, String> attributeMap = new HashMap<String, String>();
        //Populate the field mapping of UCMDB credential and vault from list data type to map data type.
        for (String attributeVault : attributesInVault) {
            attributeMap.put(attributeVault.split(":")[0], attributeVault.split(":")[1]);
        }
        //Fetch data from HashiCorp vault and update it to protocolEntry.
        populateFromVaultToCmdb(referenceID, attributeMap, protocolEntry);
        return true;
    }

    /**
     * If SSL is configured on the HashiCorp server, introduce the trust store of it.
     */
    private void initTrustStore() {
        FileInputStream input = null;
        try {
            KeyStore ks = KeyStore.getInstance("JKS");
            input = new FileInputStream(vaultFilePath + "hashiCorp.jks");
            ks.load(input, TrustStorePassword.toCharArray());
            TrustManagerFactory tmf = TrustManagerFactory.getInstance("SunX509");
            tmf.init(ks);
            TrustManager tms[] = tmf.getTrustManagers();
            /*
             * Iterate over the returned trustmanagers, look
             * for an instance of X509TrustManager.  If found,
             * use that as our "default" trust manager.
             * Customer need specify the TSL <HTTPS_VERSION>  here.
             */
            sc = SSLContext.getInstance("<HTTPS_VERSION>");
            for (int i = 0; i < tms.length; i++) {
                if (tms[i] instanceof X509TrustManager) {
                    X509TrustManager trustManager = (X509TrustManager) tms[i];
                    sc.init(null, new TrustManager[]{trustManager},
                            new java.security.SecureRandom());
                    break;
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (input != null) {
                try {
                    input.close();
                } catch (Exception e) {

                }
            }
        }
    }

    /**
     * use username and password to get token of HashiCorp, and then save the token in this provider object.
     */
    private void authenticateUserPwd() {
        DataOutputStream out = null;
        InputStream is = null;
        HttpURLConnection conn = null;
        try {
            URL url = new URL(HashiCorpAPIServer + "/auth/userpass/login/" + HashiCorpUsername);
            conn = (HttpURLConnection) url.openConnection();
            if (conn instanceof HttpsURLConnection) {
                ((HttpsURLConnection) conn).setSSLSocketFactory(sc.getSocketFactory());
                ((HttpsURLConnection) conn).setHostnameVerifier(new TrustAnyHostnameVerifier());
            }
            conn.setDoOutput(true);
            conn.setRequestMethod("POST");
            conn.setRequestProperty("Content-Type", "application/json");
            String loginJson = "{\"password\": \"" + HashiCorpPassword + "\"}";
            out = new DataOutputStream(conn.getOutputStream());
            out.write(loginJson.getBytes("UTF-8"));
            out.flush();
            is = conn.getInputStream();
            if (is != null) {
                ByteArrayOutputStream outStream = new ByteArrayOutputStream();
                byte[] buffer = new byte[1024];
                int len = 0;
                while ((len = is.read(buffer)) != -1) {
                    outStream.write(buffer, 0, len);
                }
                is.close();
                byte[] result = outStream.toByteArray();
                String response = new String(result);
                clientToken = (String) new JSONObject(response).getJSONObject("auth").get("client_token");
            }
        } catch (Exception e) {
            e.printStackTrace();
            clientToken = null;
        } finally {
            try {
                if (out != null) {
                    out.close();
                }
                if (is != null) {
                    is.close();
                }
            } catch (IOException e) {
            }
            conn.disconnect();
        }
    }

    private void authenticateByAppRole() {
        InputStream is = null;
        DataOutputStream out = null;
        HttpURLConnection conn = null;
        String roleId = null;
        String wrappedSecretId = null;
        String secretId = null;
        try {
            URL url = new URL(HashiCorpAPIServer + "/auth/approle/role/" + AppRole + "/role-id");
            conn = (HttpURLConnection) url.openConnection();
            if (conn instanceof HttpsURLConnection) {
                ((HttpsURLConnection) conn).setSSLSocketFactory(sc.getSocketFactory());
                ((HttpsURLConnection) conn).setHostnameVerifier(new TrustAnyHostnameVerifier());
            }
            conn.setDoOutput(true);
            conn.setRequestMethod("GET");
            conn.setRequestProperty("X-Vault-Token", RootToken);
            is = conn.getInputStream();
            if (is != null) {
                ByteArrayOutputStream outStream = new ByteArrayOutputStream();
                byte[] buffer = new byte[1024];
                int len = 0;
                while ((len = is.read(buffer)) != -1) {
                    outStream.write(buffer, 0, len);
                }
                byte[] result = outStream.toByteArray();
                String response = new String(result);
                roleId = (String) new JSONObject(response).getJSONObject("data").get("role_id");
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            try {
                if (is != null) {
                    is.close();
                }
            } catch (IOException e) {
            }
            conn.disconnect();
        }


        try {
            URL url = new URL(HashiCorpAPIServer + "/auth/approle/role/" + AppRole + "/secret-id");
            conn = (HttpURLConnection) url.openConnection();
            if (conn instanceof HttpsURLConnection) {
                ((HttpsURLConnection) conn).setSSLSocketFactory(sc.getSocketFactory());
                ((HttpsURLConnection) conn).setHostnameVerifier(new TrustAnyHostnameVerifier());
            }
            conn.setDoOutput(true);
            conn.setRequestMethod("POST");
            conn.setRequestProperty("X-Vault-Wrap-TTL", "600");
            conn.setRequestProperty("X-Vault-Token", RootToken);
            is = conn.getInputStream();
            if (is != null) {
                ByteArrayOutputStream outStream = new ByteArrayOutputStream();
                byte[] buffer = new byte[1024];
                int len = 0;
                while ((len = is.read(buffer)) != -1) {
                    outStream.write(buffer, 0, len);
                }
                byte[] result = outStream.toByteArray();
                String response = new String(result);
                wrappedSecretId = (String) new JSONObject(response).getJSONObject("wrap_info").get("token");
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            try {
                if (is != null) {
                    is.close();
                }
            } catch (IOException e) {
            }
            conn.disconnect();
        }


        try {
            URL url = new URL(HashiCorpAPIServer + "/sys/wrapping/unwrap");
            conn = (HttpURLConnection) url.openConnection();
            if (conn instanceof HttpsURLConnection) {
                ((HttpsURLConnection) conn).setSSLSocketFactory(sc.getSocketFactory());
                ((HttpsURLConnection) conn).setHostnameVerifier(new TrustAnyHostnameVerifier());
            }
            conn.setDoOutput(true);
            conn.setRequestMethod("POST");
            conn.setRequestProperty("X-Vault-Token", wrappedSecretId);
            is = conn.getInputStream();
            if (is != null) {
                ByteArrayOutputStream outStream = new ByteArrayOutputStream();
                byte[] buffer = new byte[1024];
                int len = 0;
                while ((len = is.read(buffer)) != -1) {
                    outStream.write(buffer, 0, len);
                }
                byte[] result = outStream.toByteArray();
                String response = new String(result);
                secretId = (String) new JSONObject(response).getJSONObject("data").get("secret_id");
            }

        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            try {
                if (is != null) {
                    is.close();
                }
            } catch (IOException e) {
            }
            conn.disconnect();
        }

        try {
            URL url = new URL(HashiCorpAPIServer + "/auth/approle/login");
            conn = (HttpURLConnection) url.openConnection();
            if (conn instanceof HttpsURLConnection) {
                ((HttpsURLConnection) conn).setSSLSocketFactory(sc.getSocketFactory());
                ((HttpsURLConnection) conn).setHostnameVerifier(new TrustAnyHostnameVerifier());
            }
            conn.setDoOutput(true);
            conn.setRequestMethod("POST");
            conn.setRequestProperty("cache-control", "no-cache");
            String loginJson = "{\"role_id\": \"" + roleId + "\" , \"secret_id\": \"" + secretId + "\"}";
            out = new DataOutputStream(conn.getOutputStream());
            out.write(loginJson.getBytes("UTF-8"));
            out.flush();
            is = conn.getInputStream();
            if (is != null) {
                ByteArrayOutputStream outStream = new ByteArrayOutputStream();
                byte[] buffer = new byte[1024];
                int len = 0;
                while ((len = is.read(buffer)) != -1) {
                    outStream.write(buffer, 0, len);
                }
                byte[] result = outStream.toByteArray();
                String response = new String(result);
                clientToken = (String) new JSONObject(response).getJSONObject("auth").get("client_token");
            }

        } catch (Exception e) {
            e.printStackTrace();
            clientToken = null;
        } finally {
            try {
                if (out != null) {
                    out.close();
                }
                if (is != null) {
                    is.close();
                }
            } catch (IOException e) {
            }
            conn.disconnect();
        }
    }

    private static class TrustAnyHostnameVerifier implements HostnameVerifier {
        public boolean verify(String hostname, SSLSession session) {
            return true;
        }
    }

    /**
     * @param referenceID   This is the reference ID in UCMDB credential and is mapped to the path in HashiCorp vault.
     * @param attributeMap  These are attributes mapped between vault and UCMDB credential, e.g. ["name":"protocol_username","pwd":"protocol_password"]
     * @param protocolEntry This the credential from UCMDB.
     */
    private void populateFromVaultToCmdb(String referenceID, Map<String, String> attributeMap, Map<String, String> protocolEntry) {
        //To avoid too frequently creating AWS iam user, provider will check lease id and use secret key and access key from local file(cache) at first.
        //If the lease id is expired, so create new access key and secret key in cache; if not, it will renew the lease id to make the access key and secret key will be continued use.
        if ("awsprotocol".equalsIgnoreCase(protocolEntry.get("protocol_type")) && awsAllAccessCache != null) {
            try {
                JSONObject awsAccessCache = (JSONObject) awsAllAccessCache.get(referenceID);
                if (awsAccessCache != null) {
                    String leaseId = awsAccessCache.get("awsLeaseId").toString();
                    if (renewLeaseId(leaseId)) {
                        protocolEntry.put("protocol_username", awsAccessCache.get("awsAccessKey").toString());
                        protocolEntry.put("protocol_password", awsAccessCache.get("awsSecretKey").toString());
                        return;
                    }
                }
            } catch (Exception e) {
                //ignore the error, as there is null when first time to access awsAllAccessCache.
            }
        }
        HttpURLConnection conn = null;
        BufferedReader br = null;
        try {
            URL url = new URL(HashiCorpAPIServer + referenceID);
            conn = (HttpURLConnection) url.openConnection();
            if (conn instanceof HttpsURLConnection) {
                ((HttpsURLConnection) conn).setSSLSocketFactory(sc.getSocketFactory());
                ((HttpsURLConnection) conn).setHostnameVerifier(new TrustAnyHostnameVerifier());
            }

            conn.setDoOutput(true);
            conn.setRequestMethod("GET");
            conn.setRequestProperty("X-Vault-Token", clientToken);
            conn.connect();
            br = new BufferedReader(new InputStreamReader(conn.getInputStream(), "UTF-8"));
            String line;
            StringBuilder sb = new StringBuilder();
            while ((line = br.readLine()) != null) {
                sb.append(line);
            }
            String response = sb.toString();
            JSONObject credentialsFromVault = new JSONObject(response).getJSONObject("data");
            if ("awsprotocol".equalsIgnoreCase(protocolEntry.get("protocol_type"))) {
                String awsSecretKey = (String) credentialsFromVault.get("secret_key");
                String awsAccessKey = (String) credentialsFromVault.get("access_key");
                String awsLeaseId = (String) new JSONObject(response).get("lease_id");
                protocolEntry.put("protocol_username", awsAccessKey);
                protocolEntry.put("protocol_password", awsSecretKey);
                //sync awsSecretKey,awsAccessKey and awsLeaseId to parameter json file with encrypted.
                JSONObject awsAccessCache = new JSONObject();
                awsAccessCache.put("awsSecretKey", awsSecretKey);
                awsAccessCache.put("awsAccessKey", awsAccessKey);
                awsAccessCache.put("awsLeaseId", awsLeaseId);
                awsAllAccessCache.put(referenceID, awsAccessCache);
                this.setParameter("awsAllAccessCache", awsAllAccessCache.toString());
                this.storeToFile(getProviderName());
            } else {
                Iterator vaultKeyIter = credentialsFromVault.keys();
                while (vaultKeyIter.hasNext()) {
                    Object vaultKey = vaultKeyIter.next();
                    String cmdbProtocolField = attributeMap.get(vaultKey);
                    if (cmdbProtocolField != null) {
                        protocolEntry.put(cmdbProtocolField, credentialsFromVault.getString((String) vaultKey));
                    }
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
            clientToken = null;
        } finally {
            try {
                if (br != null) {
                    br.close();
                }
            } catch (IOException e) {
            }
            if (conn != null) {
                conn.disconnect();
            }
        }
    }

    /**
     * check if the lease id from vault is still existing, if yes, then renew it and return true.
     *
     * @param leaseId AWS access key and secret key or STS token is bound to this lease id in vault.
     * @return true, if the lease id in vault exists; false, if the lease id in vault doesn't exist.
     */
    private boolean renewLeaseId(String leaseId) {
        if (leaseId == null) {
            return false;
        }
        HttpURLConnection conn = null;
        BufferedReader br = null;
        try {
            URL url = new URL(HashiCorpAPIServer + "/sys/leases/renew/" + leaseId);
            conn = (HttpURLConnection) url.openConnection();
            if (conn instanceof HttpsURLConnection) {
                ((HttpsURLConnection) conn).setSSLSocketFactory(sc.getSocketFactory());
                ((HttpsURLConnection) conn).setHostnameVerifier(new TrustAnyHostnameVerifier());
            }
            conn.setDoOutput(true);
            conn.setRequestMethod("PUT");
            conn.setRequestProperty("X-Vault-Token", RootToken);
            conn.connect();
            br = new BufferedReader(new InputStreamReader(conn.getInputStream(), "UTF-8"));
            String line;
            StringBuilder sb = new StringBuilder();
            while ((line = br.readLine()) != null) {
                sb.append(line);
            }
            String response = sb.toString();
            if (response.contains("lease not found")) {
                return false;
            }
        } catch (Exception e) {
            e.printStackTrace();
            return false;
        } finally {
            try {
                if (br != null) {
                    br.close();
                }
            } catch (IOException e) {
            }
            if (conn != null) {
                conn.disconnect();
            }
        }
        return true;
    }
}