// Copyright 2019 Google LLC
// Licensed 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 com.google.cloud.broker.hadoop.fs;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;

import io.grpc.ManagedChannel;
import io.grpc.Metadata;
import io.grpc.stub.MetadataUtils;
import com.google.api.client.http.GenericUrl;
import com.google.common.io.BaseEncoding;
import org.ietf.jgss.GSSException;
import org.apache.hadoop.conf.Configuration;

// Classes dynamically generated by protobuf-maven-plugin:
import com.google.cloud.broker.apps.brokerserver.protobuf.BrokerGrpc;


final class BrokerGateway {

    private BrokerGrpc.BrokerBlockingStub stub;
    private ManagedChannel managedChannel;
    private Configuration config;

    private final String CONFIG_URI = "gcp.token.broker.uri";
    private final String CONFIG_PRINCIPAL = "gcp.token.broker.kerberos.principal";
    private final String CONFIG_CERTIFICATE = "gcp.token.broker.tls.certificate";
    private final String CONFIG_CERTIFICATE_PATH = "gcp.token.broker.tls.certificate.path";


    BrokerGateway(Configuration config) {
        this(config,null);
    }

    BrokerGateway(Configuration config, String sessionToken) {
        this.config = config;

        // Extract the host and port from the URI
        String brokerUri = config.get(CONFIG_URI);
        GenericUrl url;
        try {
            url = new GenericUrl(brokerUri);
        }
        catch (IllegalArgumentException e) {
            throw new RuntimeException("Invalid value for property `" + CONFIG_URI + "`");
        }
        String host = url.getHost();
        int port = url.getPort();

        // Determine if TLS should be used
        boolean useTLS;
        String scheme = url.getScheme();
        if (scheme.equals("http")) {
            useTLS = false;
        }
        else if (scheme.equals("https")) {
            useTLS = true;
        }
        else {
            throw new RuntimeException("Incorrect URI scheme `" + scheme + " ` in `" + CONFIG_URI + "` property: " + brokerUri);
        }

        String tlsCertificate = config.get(CONFIG_CERTIFICATE);
        if (tlsCertificate == null) {
            String tlsCerfiticatePath = config.get(CONFIG_CERTIFICATE_PATH);
            if (tlsCerfiticatePath != null) {
                try {
                    tlsCertificate = new String(Files.readAllBytes(Paths.get(tlsCerfiticatePath)), StandardCharsets.US_ASCII);
                } catch (IOException e) {
                    throw new RuntimeException("Error reading the TLS certificate file: " + e.getMessage());
                }
            } else {
                tlsCertificate = "";
            }
        }

        managedChannel = GrpcUtils.newManagedChannel(host, port, useTLS, tlsCertificate);
        stub = GrpcUtils.newStub(managedChannel);

        if (sessionToken != null) {
            setSessionToken(sessionToken);
        }
        else {
            try {
                setSPNEGOToken();
            } catch (GSSException e) {
                // Clean up the channel before re-throwing the exception
                managedChannel.shutdownNow();
                throw new RuntimeException(
                    "User is not logged-in with Kerberos or cannot authenticate with the broker. Kerberos error message: " + e.getMessage());
            }
        }
    }

    BrokerGrpc.BrokerBlockingStub getStub() {
        return stub;
    }

    ManagedChannel getManagedChannel() {
        return managedChannel;
    }

    private void setSPNEGOToken() throws GSSException {
        String brokerPrincipal = config.get(CONFIG_PRINCIPAL);
        String encodedToken = BaseEncoding.base64().encode(SpnegoUtils.newSPNEGOToken(brokerPrincipal));

        // Set the 'authorization' header with the SPNEGO token
        Metadata metadata = new Metadata();
        Metadata.Key<String> key = Metadata.Key.of("authorization", Metadata.ASCII_STRING_MARSHALLER);
        metadata.put(key, "Negotiate " + encodedToken);
        stub = MetadataUtils.attachHeaders(stub, metadata);
    }

    private void setSessionToken(String sessionToken) {
        // Set the session token in the 'authorization' header
        Metadata metadata = new Metadata();
        Metadata.Key<String> key = Metadata.Key.of("authorization", Metadata.ASCII_STRING_MARSHALLER);
        metadata.put(key, "BrokerSession " + sessionToken);
        stub = MetadataUtils.attachHeaders(stub, metadata);
    }

}
@@ -636,24 +636,84 @@
 OException;%0A
+import java.net.MalformedURLException;%0Aimport java.net.URL;%0A
 import java.
@@ -894,54 +894,8 @@
 ls;%0A
-import com.google.api.client.http.GenericUrl;%0A
 impo
@@ -1899,26 +1899,19 @@
         
-GenericUrl
+URL
  url;%0A  
@@ -1948,18 +1948,11 @@
 new 
-GenericUrl
+URL
 (bro
@@ -1989,23 +1989,20 @@
 ch (
-IllegalArgument
+MalformedURL
 Exce
@@ -2267,22 +2267,24 @@
  String 
-scheme
+protocol
  = url.g
@@ -2289,14 +2289,16 @@
 .get
-Scheme
+Protocol
 ();%0A
@@ -2305,30 +2305,32 @@
         if (
-scheme
+protocol
 .equals(%22htt
@@ -2391,22 +2391,24 @@
 lse if (
-scheme
+protocol
 .equals(
@@ -2537,22 +2537,24 @@
 me %60%22 + 
-scheme
+protocol
  + %22 %60 i

Remove unnecessary dependency on GenericUrl in BrokerGateway