Credential wrapper for Domain-Wide Delegation in java using Google Cloud Libraries

DomainDelegatedCredentials.java

package com.google.auth.oauth2;

import static com.google.common.base.MoreObjects.firstNonNull;

import java.io.IOException;
import java.util.Calendar;
import java.util.Date;
import java.util.List;
import java.util.Objects;

import com.google.api.client.auth.oauth2.TokenRequest;
import com.google.api.client.auth.oauth2.TokenResponse;
import com.google.api.client.http.GenericUrl;
import com.google.api.client.googleapis.util.Utils;
import com.google.api.client.json.webtoken.JsonWebToken;
import com.google.auth.http.HttpTransportFactory;
import com.google.common.base.MoreObjects;
import com.google.api.client.http.HttpTransport;
import com.google.api.client.json.JsonFactory;

import org.apache.commons.codec.binary.Base64;

/**
 * DomainDelegatedCredentials
 * 
 */

public class DomainDelegatedCredentials extends GoogleCredentials {

  private static final long serialVersionUID = -2133257318957488431L;
  private static final String TOKEN_ENDPOINT = "https://oauth2.googleapis.com/token";

  private GoogleCredentials sourceCredentials;
  private List<String> scopes;
  private String subject;
  private final String transportFactoryClassName;
  private transient HttpTransportFactory transportFactory;

  public static GoogleCredentials create(GoogleCredentials sourceCredentials, String subject,
      List<String> scopes) {
    return DomainDelegatedCredentials.newBuilder().setSourceCredentials(sourceCredentials).setScopes(scopes)
        .setSubject(subject).build();
  }

  public static GoogleCredentials create(GoogleCredentials sourceCredentials, String subject,
      List<String> scopes, HttpTransportFactory transportFactory) {
    return DomainDelegatedCredentials.newBuilder().setSourceCredentials(sourceCredentials).setSubject(subject)
        .setScopes(scopes).setHttpTransportFactory(transportFactory).build();
  }

  private DomainDelegatedCredentials(Builder builder) {
    this.sourceCredentials = builder.getSourceCredentials();
    this.scopes = builder.getScopes();
    this.subject = builder.getSubject();
    this.transportFactory = firstNonNull(builder.getHttpTransportFactory(),
        getFromServiceLoader(HttpTransportFactory.class, OAuth2Utils.HTTP_TRANSPORT_FACTORY));
    this.transportFactoryClassName = this.transportFactory.getClass().getName();
  }

  @Override
  public AccessToken refreshAccessToken() throws IOException {

    if (this.sourceCredentials.createScopedRequired()) {
      this.sourceCredentials = this.sourceCredentials.createScoped(this.scopes);
    }

    try {
      this.sourceCredentials.refreshIfExpired();
    } catch (IOException e) {
      throw new IOException("Unable to refresh sourceCredentials", e);
    }

    // https://googleapis.dev/java/google-auth-library/latest/index.html?com/google/auth/oauth2/GoogleCredentials.html
    /*
     * All Known Implementing Classes: AppEngineCredentials,
     * ComputeEngineCredentials, ImpersonatedCredentials, ServiceAccountCredentials,
     **** ServiceAccountJwtAccessCredentials
     */

    String iss = "";

    if (this.sourceCredentials instanceof ServiceAccountCredentials)
      iss = ((ServiceAccountCredentials) this.sourceCredentials).getAccount();
    else if (this.sourceCredentials instanceof ComputeEngineCredentials)
      iss = ((ComputeEngineCredentials) this.sourceCredentials).getAccount();
    else if (this.sourceCredentials instanceof ImpersonatedCredentials)
      iss = ((ImpersonatedCredentials) this.sourceCredentials).getAccount();
    else if (this.sourceCredentials instanceof AppEngineCredentials)
      iss = ((AppEngineCredentials) this.sourceCredentials).getAccount();
    else
      throw new IOException("Provided source credential does not implement ServiceAccountSigner");

    long now = System.currentTimeMillis() / 1000L;
    long exp = now + 30;

    HttpTransport httpTransport = Utils.getDefaultTransport();
    JsonFactory jsonFactory = Utils.getDefaultJsonFactory();

    JsonWebToken.Header header = new JsonWebToken.Header();
    header.setFactory(jsonFactory);
    header.setType("JWT");
    header.set("alg", "RS256");

    JsonWebToken.Payload payload = new JsonWebToken.Payload();
    payload.setFactory(jsonFactory);
    payload.setAudience(TOKEN_ENDPOINT);

    StringBuffer sb = new StringBuffer();
    for (String s : this.scopes) {
      sb.append(s);
      sb.append(" ");
    }

    payload.set("scope", sb.toString());
    payload.setIssuedAtTimeSeconds(now);
    payload.setExpirationTimeSeconds(exp);
    payload.setIssuer(iss);
    payload.setSubject(this.subject);

    String jwt = Base64.encodeBase64URLSafeString(header.toPrettyString().getBytes()) + "."
        + Base64.encodeBase64URLSafeString(payload.toPrettyString().getBytes("UTF-8"));

    byte[] signatureBytes;

    if (this.sourceCredentials instanceof ServiceAccountCredentials)
      signatureBytes = ((ServiceAccountCredentials) this.sourceCredentials).sign(jwt.getBytes("UTF-8"));
    else if (this.sourceCredentials instanceof ComputeEngineCredentials)
      signatureBytes = ((ComputeEngineCredentials) this.sourceCredentials).sign(jwt.getBytes("UTF-8"));
    else if (this.sourceCredentials instanceof ImpersonatedCredentials)
      signatureBytes = ((ImpersonatedCredentials) this.sourceCredentials).sign(jwt.getBytes("UTF-8"));
    else if (this.sourceCredentials instanceof AppEngineCredentials)
      signatureBytes = ((AppEngineCredentials) this.sourceCredentials).sign(jwt.getBytes("UTF-8"));
    else
      throw new IOException("Provided source credential does not implement ServiceAccountSigner");

    String b64sig = Base64.encodeBase64URLSafeString(signatureBytes);
    String assertion = jwt + "." + b64sig;

    TokenRequest request = new TokenRequest(this.transportFactory.create(), jsonFactory,
        new GenericUrl(TOKEN_ENDPOINT), "assertion");
    request.put("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer");
    request.put("assertion", assertion);
    TokenResponse response = request.execute();

    Date rnow = new Date();
    Calendar cal = Calendar.getInstance();
    cal.setTime(rnow);

    cal.add(Calendar.SECOND, Math.toIntExact(response.getExpiresInSeconds()));
    return new AccessToken(response.getAccessToken(), cal.getTime());
  }

  @Override
  public int hashCode() {
    return Objects.hash(sourceCredentials, scopes, subject);
  }

  @Override
  public String toString() {
    return MoreObjects.toStringHelper(this).add("sourceCredentials", sourceCredentials).add("scopes", scopes)
        .add("subject", subject).add("transportFactoryClassName", transportFactoryClassName).toString();
  }

  @Override
  public boolean equals(Object obj) {
    if (!(obj instanceof DomainDelegatedCredentials)) {
      return false;
    }
    DomainDelegatedCredentials other = (DomainDelegatedCredentials) obj;
    return Objects.equals(this.sourceCredentials, other.sourceCredentials)
        && Objects.equals(this.subject, other.subject) && Objects.equals(this.scopes, other.scopes)
        && Objects.equals(this.transportFactoryClassName, other.transportFactoryClassName);
  }

  public Builder toBuilder() {
    return new Builder(this.sourceCredentials, this.subject, this.scopes);
  }

  public static Builder newBuilder() {
    return new Builder();
  }

  public static class Builder extends GoogleCredentials.Builder {

    private GoogleCredentials sourceCredentials;
    private List<String> scopes;
    private String subject;
    private HttpTransportFactory transportFactory;

    protected Builder() {
    }

    protected Builder(GoogleCredentials sourceCredentials, String subject, List<String> scopes) {
      this.sourceCredentials = sourceCredentials;
      this.subject = subject;
      this.scopes = scopes;
    }

    public Builder setSourceCredentials(GoogleCredentials sourceCredentials) {
      this.sourceCredentials = sourceCredentials;
      return this;
    }

    public GoogleCredentials getSourceCredentials() {
      return this.sourceCredentials;
    }

    public Builder setScopes(List<String> scopes) {
      this.scopes = scopes;
      return this;
    }

    public List<String> getScopes() {
      return this.scopes;
    }

    public Builder setSubject(String subject) {
      this.subject = subject;
      return this;
    }

    public String getSubject() {
      return this.subject;
    }

    public Builder setHttpTransportFactory(HttpTransportFactory transportFactory) {
      this.transportFactory = transportFactory;
      return this;
    }

    public HttpTransportFactory getHttpTransportFactory() {
      return transportFactory;
    }

    public DomainDelegatedCredentials build() {

      if (!(this.sourceCredentials instanceof ServiceAccountCredentials
          || this.sourceCredentials instanceof ComputeEngineCredentials
          || this.sourceCredentials instanceof ImpersonatedCredentials
          || this.sourceCredentials instanceof AppEngineCredentials)) {

      }

      return new DomainDelegatedCredentials(this);
    }
  }
}