<commit_before>package com.emc.ecs.management.sdk;

import com.emc.ecs.management.sdk.model.EcsManagementClientError;
import com.emc.ecs.servicebroker.EcsManagementClientException;
import com.emc.ecs.servicebroker.EcsManagementClientUnauthorizedException;
import com.emc.ecs.servicebroker.EcsManagementResourceNotFoundException;
import org.glassfish.jersey.client.authentication.HttpAuthenticationFeature;
import org.glassfish.jersey.logging.LoggingFeature;
import org.slf4j.LoggerFactory;

import javax.net.ssl.HostnameVerifier;
import javax.net.ssl.HttpsURLConnection;
import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManagerFactory;
import javax.ws.rs.client.Client;
import javax.ws.rs.client.ClientBuilder;
import javax.ws.rs.client.Entity;
import javax.ws.rs.client.Invocation.Builder;
import javax.ws.rs.core.Response;
import javax.ws.rs.core.UriBuilder;
import java.io.ByteArrayInputStream;
import java.io.InputStream;
import java.security.KeyStore;
import java.security.cert.Certificate;
import java.security.cert.CertificateFactory;
import java.util.logging.Level;
import java.util.logging.Logger;

import static com.emc.ecs.management.sdk.Constants.*;

public class Connection {
    private static final org.slf4j.Logger logger = LoggerFactory.getLogger(Connection.class);

    private static final int AUTH_RETRIES_MAX = 3;
    private final String endpoint;
    private final String username;
    private final String password;
    private String authToken;
    private String certificate;
    private int authRetries = 0;

    public Connection(String endpoint, String username, String password) {
        super();
        this.endpoint = endpoint;
        this.username = username;
        this.password = password;
    }

    public Connection(String endpoint, String username, String password,
                      String certificate) {
        super();
        this.endpoint = endpoint;
        this.username = username;
        this.password = password;
        this.certificate = certificate;
    }

    private static HostnameVerifier getHostnameVerifier() {
        return (hostname, session) -> true;
    }

    public String getAuthToken() {
        return authToken;
    }

    private Client buildJerseyClient() throws EcsManagementClientException {
        ClientBuilder builder;
        if (certificate != null) {
            // Disable host name verification. Should be able to configure the
            // ECS certificate with the correct host name to avoid this.
            HostnameVerifier hostnameVerifier = getHostnameVerifier();
            HttpsURLConnection.setDefaultHostnameVerifier(hostnameVerifier);
            builder = ClientBuilder.newBuilder()
                    .register(hostnameVerifier);
            builder.sslContext(getSSLContext());
        } else {
            builder = ClientBuilder.newBuilder();
        }
        return builder.build();
    }

    private SSLContext getSSLContext() throws EcsManagementClientException {
        try {
            CertificateFactory certFactory;

            //InputStream targetStream = new ByteArrayInputStream(initialString.getBytes());
            //InputStream certInputStream = certificate.openStream();

            InputStream certInputStream = new ByteArrayInputStream(certificate.getBytes());

            SSLContext sslContext = SSLContext.getInstance("TLS");
            certFactory = CertificateFactory.getInstance("X.509");
            Certificate caCert = certFactory
                    .generateCertificate(certInputStream);
            TrustManagerFactory trustMgrFactory = TrustManagerFactory
                    .getInstance(TrustManagerFactory.getDefaultAlgorithm());
            KeyStore keyStore = KeyStore.getInstance(KeyStore.getDefaultType());
            keyStore.load(null);
            keyStore.setCertificateEntry("caCert", caCert);
            trustMgrFactory.init(keyStore);
            sslContext.init(null, trustMgrFactory.getTrustManagers(), null);
            return sslContext;
        } catch (Exception e) {
            throw new EcsManagementClientException(e);
        }
    }

    public boolean isLoggedIn() {
        return authToken != null;
    }

    public void login() throws EcsManagementClientException {
        UriBuilder uriBuilder = UriBuilder.fromPath(endpoint).segment("login");

        logger.info("Logging into {} as {}", endpoint, username);

        HttpAuthenticationFeature authFeature = HttpAuthenticationFeature
                .basicBuilder().credentials(username, password).build();
        Client jerseyClient = buildJerseyClient().register(authFeature);

        Builder request = jerseyClient.target(uriBuilder).request();

        Response response = null;

        try {
            response = request.get();
            handleResponse(response);

            this.authToken = response.getHeaderString("X-SDS-AUTH-TOKEN");
            this.authRetries = 0;
        } catch (EcsManagementResourceNotFoundException e) {
            logger.warn("Login failed to handle response: {}", e.getMessage());
            logger.warn("Response: {}", response);

            throw new EcsManagementClientException(e);
        } catch (EcsManagementClientUnauthorizedException e) {
            logger.error("Failed to login to ECS management endpoint {}: {}", endpoint, e.getMessage());
            throw new EcsManagementClientUnauthorizedException("Login attempt failed: " + e.toString(), e);
        } catch (Exception e) {
            logger.error("Failed to login to ECS management endpoint {}: {}", endpoint, e.getMessage());
            throw new EcsManagementClientException(e.getMessage(), e);
        }
    }

    public void logout() throws EcsManagementClientException {
        UriBuilder uri = UriBuilder.fromPath(endpoint).segment("logout")
                .queryParam("force", true);
        handleRemoteCall(GET, uri, null);
        this.authToken = null;
    }

    protected Response handleRemoteCall(String method, UriBuilder uri, Object arg) throws EcsManagementClientException {
        return handleRemoteCall(method, uri, arg, XML);
    }

    protected Response handleRemoteCall(String method, UriBuilder uri,
                                        Object arg, String contentType) throws EcsManagementClientException {
        Response response = makeRemoteCall(method, uri, arg, contentType);
        try {
            handleResponse(response);
        } catch (EcsManagementResourceNotFoundException e) {
            throw new EcsManagementClientException(e);
        }
        return response;
    }

    protected UriBuilder getUriBuilder() {
        return UriBuilder.fromPath(endpoint);
    }

    protected Response makeRemoteCall(String method, UriBuilder uri, Object arg) throws EcsManagementClientException {
        return makeRemoteCall(method, uri, arg, XML);
    }

    protected Response makeRemoteCall(String method, UriBuilder uri, Object arg, String contentType)
            throws EcsManagementClientException {
        if (!isLoggedIn()) {
            login();
        }

        try {
            logger.info("{} {}", method, uri);

            Client jerseyClient = buildJerseyClient();
            Builder request = jerseyClient.target(uri)
                    .register(LoggingFeature.class).request()
                    .header("X-SDS-AUTH-TOKEN", authToken)
                    .header("Accept", "application/xml");

            Response response = null;
            if (GET.equals(method)) {
                response = request.get();
            } else if (POST.equals(method) || PUT.equals(method)) {
                Entity<Object> objectEntity;
                if (XML.equals(contentType)) {
                    objectEntity = Entity.xml(arg);<commit_msg>Add override header to to support ECS Flex (#164)<commit_after>package com.emc.ecs.management.sdk;

import com.emc.ecs.management.sdk.model.EcsManagementClientError;
import com.emc.ecs.servicebroker.EcsManagementClientException;
import com.emc.ecs.servicebroker.EcsManagementClientUnauthorizedException;
import com.emc.ecs.servicebroker.EcsManagementResourceNotFoundException;
import org.glassfish.jersey.client.authentication.HttpAuthenticationFeature;
import org.glassfish.jersey.logging.LoggingFeature;
import org.slf4j.LoggerFactory;

import javax.net.ssl.HostnameVerifier;
import javax.net.ssl.HttpsURLConnection;
import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManagerFactory;
import javax.ws.rs.client.Client;
import javax.ws.rs.client.ClientBuilder;
import javax.ws.rs.client.Entity;
import javax.ws.rs.client.Invocation.Builder;
import javax.ws.rs.core.Response;
import javax.ws.rs.core.UriBuilder;
import java.io.ByteArrayInputStream;
import java.io.InputStream;
import java.security.KeyStore;
import java.security.cert.Certificate;
import java.security.cert.CertificateFactory;
import java.util.logging.Level;
import java.util.logging.Logger;

import static com.emc.ecs.management.sdk.Constants.*;

public class Connection {
    private static final org.slf4j.Logger logger = LoggerFactory.getLogger(Connection.class);

    private static final int AUTH_RETRIES_MAX = 3;
    private final String endpoint;
    private final String username;
    private final String password;
    private String authToken;
    private String certificate;
    private int authRetries = 0;

    public Connection(String endpoint, String username, String password) {
        super();
        this.endpoint = endpoint;
        this.username = username;
        this.password = password;
    }

    public Connection(String endpoint, String username, String password,
                      String certificate) {
        super();
        this.endpoint = endpoint;
        this.username = username;
        this.password = password;
        this.certificate = certificate;
    }

    private static HostnameVerifier getHostnameVerifier() {
        return (hostname, session) -> true;
    }

    public String getAuthToken() {
        return authToken;
    }

    private Client buildJerseyClient() throws EcsManagementClientException {
        ClientBuilder builder;
        if (certificate != null) {
            // Disable host name verification. Should be able to configure the
            // ECS certificate with the correct host name to avoid this.
            HostnameVerifier hostnameVerifier = getHostnameVerifier();
            HttpsURLConnection.setDefaultHostnameVerifier(hostnameVerifier);
            builder = ClientBuilder.newBuilder()
                    .register(hostnameVerifier);
            builder.sslContext(getSSLContext());
        } else {
            builder = ClientBuilder.newBuilder();
        }
        return builder.build();
    }

    private SSLContext getSSLContext() throws EcsManagementClientException {
        try {
            CertificateFactory certFactory;

            //InputStream targetStream = new ByteArrayInputStream(initialString.getBytes());
            //InputStream certInputStream = certificate.openStream();

            InputStream certInputStream = new ByteArrayInputStream(certificate.getBytes());

            SSLContext sslContext = SSLContext.getInstance("TLS");
            certFactory = CertificateFactory.getInstance("X.509");
            Certificate caCert = certFactory
                    .generateCertificate(certInputStream);
            TrustManagerFactory trustMgrFactory = TrustManagerFactory
                    .getInstance(TrustManagerFactory.getDefaultAlgorithm());
            KeyStore keyStore = KeyStore.getInstance(KeyStore.getDefaultType());
            keyStore.load(null);
            keyStore.setCertificateEntry("caCert", caCert);
            trustMgrFactory.init(keyStore);
            sslContext.init(null, trustMgrFactory.getTrustManagers(), null);
            return sslContext;
        } catch (Exception e) {
            throw new EcsManagementClientException(e);
        }
    }

    public boolean isLoggedIn() {
        return authToken != null;
    }

    public void login() throws EcsManagementClientException {
        UriBuilder uriBuilder = UriBuilder.fromPath(endpoint).segment("login");

        logger.info("Logging into {} as {}", endpoint, username);

        HttpAuthenticationFeature authFeature = HttpAuthenticationFeature
                .basicBuilder().credentials(username, password).build();
        Client jerseyClient = buildJerseyClient().register(authFeature);

        Builder request = jerseyClient.target(uriBuilder).request();

        Response response = null;

        try {
            response = request.get();
            handleResponse(response);

            this.authToken = response.getHeaderString("X-SDS-AUTH-TOKEN");
            this.authRetries = 0;
        } catch (EcsManagementResourceNotFoundException e) {
            logger.warn("Login failed to handle response: {}", e.getMessage());
            logger.warn("Response: {}", response);

            throw new EcsManagementClientException(e);
        } catch (EcsManagementClientUnauthorizedException e) {
            logger.error("Failed to login to ECS management endpoint {}: {}", endpoint, e.getMessage());
            throw new EcsManagementClientUnauthorizedException("Login attempt failed: " + e.toString(), e);
        } catch (Exception e) {
            logger.error("Failed to login to ECS management endpoint {}: {}", endpoint, e.getMessage());
            throw new EcsManagementClientException(e.getMessage(), e);
        }
    }

    public void logout() throws EcsManagementClientException {
        UriBuilder uri = UriBuilder.fromPath(endpoint).segment("logout")
                .queryParam("force", true);
        handleRemoteCall(GET, uri, null);
        this.authToken = null;
    }

    protected Response handleRemoteCall(String method, UriBuilder uri, Object arg) throws EcsManagementClientException {
        return handleRemoteCall(method, uri, arg, XML);
    }

    protected Response handleRemoteCall(String method, UriBuilder uri,
                                        Object arg, String contentType) throws EcsManagementClientException {
        Response response = makeRemoteCall(method, uri, arg, contentType);
        try {
            handleResponse(response);
        } catch (EcsManagementResourceNotFoundException e) {
            throw new EcsManagementClientException(e);
        }
        return response;
    }

    protected UriBuilder getUriBuilder() {
        return UriBuilder.fromPath(endpoint);
    }

    protected Response makeRemoteCall(String method, UriBuilder uri, Object arg) throws EcsManagementClientException {
        return makeRemoteCall(method, uri, arg, XML);
    }

    protected Response makeRemoteCall(String method, UriBuilder uri, Object arg, String contentType)
            throws EcsManagementClientException {
        if (!isLoggedIn()) {
            login();
        }

        try {
            logger.info("{} {}", method, uri);

            Client jerseyClient = buildJerseyClient();
            Builder request = jerseyClient.target(uri)
                    .register(LoggingFeature.class).request()
                    .header("X-EMC-Override", "true")            // enables access to ECS Flex API (pre-GA limitation)
                    .header("X-SDS-AUTH-TOKEN", authToken)
                    .header("Accept", "application/xml");

            Response response = null;
            if (GET.equals(method)) {
                response = request.get();
            } else if (POST.equals(method) || PUT.equals(method)) {
                Entity<Object> objectEntity;
                if (XML.equals(contentType)) {
                    objectEntity = Entity.xml(arg);
{'fraction_non_alphanumeric': 0.09132012991148188, 'fraction_numerical': 0.001209959880277654, 'mean_word_length': 2.7488660778228695, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 4}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2492893', 'n_tokens_mistral': 3898, 'n_tokens_neox': 3788, 'n_words': 903}