Add getIdentityCertificate and getRoleCertificate to ZtsClient

4 files changed, 122 insertions, 6 deletions

diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/impl/AthenzClientFactoryImpl.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/impl/AthenzClientFactoryImpl.java
index 540cdb39630..e640fcd0ea7 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/impl/AthenzClientFactoryImpl.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/impl/AthenzClientFactoryImpl.java
@@ -51,7 +51,7 @@ public class AthenzClientFactoryImpl implements AthenzClientFactory {
      */
     @Override
     public ZtsClient createZtsClientWithServicePrincipal() {
-        return new ZtsClientImpl(new ZTSClient(config.ztsUrl(), createServicePrincipal()), config);
+        return new ZtsClientImpl(new ZTSClient(config.ztsUrl(), createServicePrincipal()), getServicePrivateKey(), config);
     }
 
     /**
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/impl/ZtsClientImpl.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/impl/ZtsClientImpl.java
index 76c000936e3..d518e1905e5 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/impl/ZtsClientImpl.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/impl/ZtsClientImpl.java
@@ -1,18 +1,27 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.controller.athenz.impl;
 
+import com.yahoo.athenz.auth.util.Crypto;
+import com.yahoo.athenz.zts.InstanceRefreshRequest;
+import com.yahoo.athenz.zts.RoleCertificateRequest;
 import com.yahoo.athenz.zts.TenantDomains;
 import com.yahoo.athenz.zts.ZTSClient;
 import com.yahoo.athenz.zts.ZTSClientException;
 import com.yahoo.log.LogLevel;
 import com.yahoo.vespa.hosted.controller.api.identifiers.AthenzDomain;
 import com.yahoo.vespa.hosted.controller.api.integration.athenz.AthenzIdentity;
+import com.yahoo.vespa.hosted.controller.api.integration.athenz.AthenzIdentityCertificate;
+import com.yahoo.vespa.hosted.controller.api.integration.athenz.AthenzRoleCertificate;
 import com.yahoo.vespa.hosted.controller.api.integration.athenz.AthenzService;
 import com.yahoo.vespa.hosted.controller.api.integration.athenz.ZtsClient;
 import com.yahoo.vespa.hosted.controller.api.integration.athenz.ZtsException;
 import com.yahoo.vespa.hosted.controller.athenz.config.AthenzConfig;
 
+import java.security.PrivateKey;
+import java.security.cert.X509Certificate;
+import java.time.Duration;
 import java.util.List;
+import java.util.function.Supplier;
 import java.util.logging.Logger;
 
 import static java.util.stream.Collectors.toList;
@@ -23,27 +32,80 @@ import static java.util.stream.Collectors.toList;
 public class ZtsClientImpl implements ZtsClient {
 
     private static final Logger log = Logger.getLogger(ZtsClientImpl.class.getName());
+    private static final Duration CERTIFICATE_EXPIRY = Duration.ofHours(1);
 
     private final ZTSClient ztsClient;
     private final AthenzService service;
+    private final PrivateKey privateKey;
+    private final String certificateDnsDomain;
 
-    public ZtsClientImpl(ZTSClient ztsClient, AthenzConfig config) {
+    public ZtsClientImpl(ZTSClient ztsClient, PrivateKey privateKey, AthenzConfig config) {
         this.ztsClient = ztsClient;
         this.service = new AthenzService(config.domain(), config.service().name());
+        this.privateKey = privateKey;
+        this.certificateDnsDomain = config.certDnsDomain();
     }
 
     @Override
     public List<AthenzDomain> getTenantDomainsForUser(AthenzIdentity identity) {
-        log.log(LogLevel.DEBUG, String.format(
-                "getTenantDomains(domain=%s, identity=%s, rolename=admin, service=%s)",
-                service.getDomain().id(), identity.getFullName(), service.getFullName()));
-        try {
+        return getOrThrow(() -> {
+            log.log(LogLevel.DEBUG, String.format(
+                    "getTenantDomains(domain=%s, identity=%s, rolename=admin, service=%s)",
+                    service.getDomain().id(), identity.getFullName(), service.getFullName()));
             TenantDomains domains = ztsClient.getTenantDomains(
                     service.getDomain().id(), identity.getFullName(), "admin", service.getName());
             return domains.getTenantDomainNames().stream()
                     .map(AthenzDomain::new)
                     .collect(toList());
+        });
+    }
+
+    @Override
+    public AthenzIdentityCertificate getIdentityCertificate() {
+        return getOrThrow(() -> {
+            log.log(LogLevel.DEBUG,
+                    String.format("postInstanceRefreshRequest(service=%s)", service.getFullName()));
+            InstanceRefreshRequest req =
+                    ZTSClient.generateInstanceRefreshRequest(
+                            service.getDomain().id(),
+                            service.getName(),
+                            privateKey,
+                            certificateDnsDomain,
+                            (int) CERTIFICATE_EXPIRY.getSeconds());
+            X509Certificate certificate = Crypto.loadX509Certificate(
+                    ztsClient.postInstanceRefreshRequest(service.getDomain().id(), service.getName(), req)
+                            .getCertificate());
+            return new AthenzIdentityCertificate(certificate, privateKey);
+        });
+    }
+
+    @Override
+    public AthenzRoleCertificate getRoleCertificate(AthenzDomain roleDomain, String roleName) {
+        return getOrThrow(() -> {
+            log.log(LogLevel.DEBUG,
+                    String.format("postRoleCertificateRequest(service=%s, roleDomain=%s, roleName=%s)",
+                                  service.getFullName(), roleDomain.id(), roleName));
+            RoleCertificateRequest req =
+                    ZTSClient.generateRoleCertificateRequest(
+                            service.getDomain().id(),
+                            service.getName(),
+                            roleDomain.id(),
+                            roleName,
+                            privateKey,
+                            certificateDnsDomain,
+                            (int)CERTIFICATE_EXPIRY.getSeconds());
+            X509Certificate roleCertificate = Crypto.loadX509Certificate(
+                    ztsClient.postRoleCertificateRequest(roleDomain.id(), roleName, req)
+                            .getToken());
+            return new AthenzRoleCertificate(roleCertificate, privateKey);
+        });
+    }
+
+    private static <T> T getOrThrow(Supplier<T> wrappedCode) {
+        try {
+            return wrappedCode.get();
         } catch (ZTSClientException e) {
+            log.warning("Error from Athenz: " + e.getMessage());
             throw new ZtsException(e.getCode(), e);
         }
     }
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/mock/ZtsClientMock.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/mock/ZtsClientMock.java
index ee43181a358..d778fb550ed 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/mock/ZtsClientMock.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/athenz/mock/ZtsClientMock.java
@@ -1,10 +1,21 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.controller.athenz.mock;
 
+import com.yahoo.athenz.auth.util.Crypto;
 import com.yahoo.vespa.hosted.controller.api.identifiers.AthenzDomain;
 import com.yahoo.vespa.hosted.controller.api.integration.athenz.AthenzIdentity;
+import com.yahoo.vespa.hosted.controller.api.integration.athenz.AthenzIdentityCertificate;
+import com.yahoo.vespa.hosted.controller.api.integration.athenz.AthenzRoleCertificate;
 import com.yahoo.vespa.hosted.controller.api.integration.athenz.ZtsClient;
+import org.bouncycastle.asn1.x500.X500Name;
+import org.bouncycastle.operator.OperatorCreationException;
+import org.bouncycastle.pkcs.PKCS10CertificationRequest;
 
+import java.io.IOException;
+import java.security.KeyPair;
+import java.security.KeyPairGenerator;
+import java.security.NoSuchAlgorithmException;
+import java.security.cert.X509Certificate;
 import java.util.List;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -31,4 +42,44 @@ public class ZtsClientMock implements ZtsClient {
                 .map(domain -> domain.name)
                 .collect(toList());
     }
+
+    @Override
+    public AthenzIdentityCertificate getIdentityCertificate() {
+        log.log(Level.INFO, "getIdentityCertificate()");
+        try {
+            KeyPair keyPair = createKeyPair();
+            String subject = "CN=controller";
+            return new AthenzIdentityCertificate(createCertificate(keyPair, subject), keyPair.getPrivate());
+        } catch (NoSuchAlgorithmException | OperatorCreationException | IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public AthenzRoleCertificate getRoleCertificate(AthenzDomain roleDomain, String roleName) {
+        log.log(Level.INFO,
+                String.format("getRoleCertificate(roleDomain=%s, roleName=%s)", roleDomain.id(), roleDomain));
+        try {
+            KeyPair keyPair = createKeyPair();
+            String subject = String.format("CN=%s:role.%s", roleDomain.id(), roleName);
+            return new AthenzRoleCertificate(createCertificate(keyPair, subject), keyPair.getPrivate());
+        } catch (NoSuchAlgorithmException | OperatorCreationException | IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private static X509Certificate createCertificate(KeyPair keyPair, String subject) throws
+            OperatorCreationException, IOException {
+        PKCS10CertificationRequest csr =
+                Crypto.getPKCS10CertRequest(
+                        Crypto.generateX509CSR(keyPair.getPrivate(), subject, null));
+        return Crypto.generateX509Certificate(csr, keyPair.getPrivate(), new X500Name(subject), 3600, false);
+    }
+
+    private static KeyPair createKeyPair() throws NoSuchAlgorithmException {
+        KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA");
+        keyGen.initialize(512);
+        return keyGen.genKeyPair();
+    }
+
 }
diff --git a/controller-server/src/main/resources/configdefinitions/athenz.def b/controller-server/src/main/resources/configdefinitions/athenz.def
index 6d10f3dee28..2ee15e183f2 100644
--- a/controller-server/src/main/resources/configdefinitions/athenz.def
+++ b/controller-server/src/main/resources/configdefinitions/athenz.def
@@ -17,6 +17,9 @@ domain                          string
 userAuthenticationPassThruAttribute  string
 # TODO Remove once migrated to Okta
 
+# Certificate DNS domain
+certDnsDomain                   string
+
 # Athenz service name for controller identity
 service.name                    string