// 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.google.inject.Inject;
import com.yahoo.athenz.auth.Principal;
import com.yahoo.athenz.auth.impl.PrincipalAuthority;
import com.yahoo.athenz.auth.impl.SimplePrincipal;
import com.yahoo.athenz.auth.impl.SimpleServiceIdentityProvider;
import com.yahoo.athenz.auth.token.PrincipalToken;
import com.yahoo.athenz.auth.util.Crypto;
import com.yahoo.athenz.zms.ZMSClient;
import com.yahoo.athenz.zts.ZTSClient;
import com.yahoo.vespa.athenz.api.NToken;
import com.yahoo.vespa.athenz.utils.AthenzIdentities;
import com.yahoo.vespa.hosted.controller.api.integration.athenz.AthenzClientFactory;
import com.yahoo.vespa.hosted.controller.api.integration.athenz.ZmsClient;
import com.yahoo.vespa.hosted.controller.api.integration.athenz.ZtsClient;
import com.yahoo.vespa.hosted.controller.api.integration.security.KeyService;
import com.yahoo.vespa.hosted.controller.athenz.config.AthenzConfig;

import java.security.PrivateKey;
import java.time.Duration;

/**
 * @author bjorncs
 */
public class AthenzClientFactoryImpl implements AthenzClientFactory {

    private final KeyService secretService;
    private final AthenzConfig config;
    private final AthenzPrincipalAuthority athenzPrincipalAuthority;

    @Inject
    public AthenzClientFactoryImpl(KeyService secretService, AthenzConfig config) {
        this.secretService = secretService;
        this.config = config;
        this.athenzPrincipalAuthority = new AthenzPrincipalAuthority(config.principalHeaderName());
    }

    /**
     * @return A ZMS client instance with the service identity as principal.
     */
    @Override
    public ZmsClient createZmsClientWithServicePrincipal() {
        return new ZmsClientImpl(new ZMSClient(config.zmsUrl(), createServicePrincipal()), config);
    }

    /**
     * @return A ZTS client instance with the service identity as principal.
     */
    @Override
    public ZtsClient createZtsClientWithServicePrincipal() {
        return new ZtsClientImpl(new ZTSClient(config.ztsUrl(), createServicePrincipal()), getServicePrivateKey(), config);
    }

    /**
     * @return A ZMS client created with a dual principal representing both the tenant admin and the service identity.
     */
    @Override
    public ZmsClient createZmsClientWithAuthorizedServiceToken(NToken authorizedServiceToken) {
        PrincipalToken signedToken = new PrincipalToken(authorizedServiceToken.getRawToken());
        AthenzConfig.Service service = config.service();
        signedToken.signForAuthorizedService(
                config.domain() + "." + service.name(), service.publicKeyId(), getServicePrivateKey());

        Principal dualPrincipal = SimplePrincipal.create(
                AthenzIdentities.USER_PRINCIPAL_DOMAIN.getName(), signedToken.getName(), signedToken.getSignedToken(), athenzPrincipalAuthority);
        return new ZmsClientImpl(new ZMSClient(config.zmsUrl(), dualPrincipal), config);

    }

    private Principal createServicePrincipal() {
        AthenzConfig.Service service = config.service();
        // TODO bjorncs: Cache principal token
        SimpleServiceIdentityProvider identityProvider =
                new SimpleServiceIdentityProvider(
                        athenzPrincipalAuthority,
                        config.domain(),
                        service.name(),
                        getServicePrivateKey(),
                        service.publicKeyId(),
                        Duration.ofMinutes(service.credentialsExpiryMinutes()).getSeconds());
        return identityProvider.getIdentity(config.domain(), service.name());
    }

    private PrivateKey getServicePrivateKey() {
        AthenzConfig.Service service = config.service();
        String privateKey = secretService.getSecret(service.privateKeySecretName(), service.privateKeyVersion()).trim();
        return Crypto.loadPrivateKey(privateKey);
    }

    private static class AthenzPrincipalAuthority extends PrincipalAuthority {
        private final String principalHeaderName;

        public AthenzPrincipalAuthority(String principalHeaderName) {
            this.principalHeaderName = principalHeaderName;
        }

        @Override
        public String getHeader() {
            return principalHeaderName;
        }
    }


}

@@ -574,32 +574,71 @@
 .zts.ZTSClient;%0A
+import com.yahoo.container.jdisc.Ckms;%0A
 import com.yahoo
@@ -659,24 +659,24 @@
 api.NToken;%0A
-
 import com.y
@@ -956,86 +956,8 @@
 nt;%0A
-import com.yahoo.vespa.hosted.controller.api.integration.security.KeyService;%0A
 impo
@@ -1110,16 +1110,48 @@
 ncs%0A */%0A
+// TODO Use SiaIdentityProvider%0A
 public c
@@ -1235,32 +1235,17 @@
 nal 
-KeyService secretService
+Ckms ckms
 ;%0A  
@@ -1402,32 +1402,17 @@
 mpl(
-KeyService secretService
+Ckms ckms
 , At
@@ -1449,37 +1449,19 @@
 his.
-secretService = secretService
+ckms = ckms
 ;%0A  
@@ -3821,21 +3821,12 @@
 y = 
-secretService
+ckms
 .get

Use CKMS instead of KeyService to retrieve service private key