/*
 * Copyright 2016 the original author or authors.
 *
 * 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 org.springframework.cloud.vault.util;

import java.util.Collections;
import java.util.Map;

import org.springframework.util.Assert;
import org.springframework.util.StringUtils;
import org.springframework.vault.core.VaultOperations;
import org.springframework.vault.core.VaultSysOperations;
import org.springframework.vault.support.VaultHealth;
import org.springframework.vault.support.VaultInitializationRequest;
import org.springframework.vault.support.VaultInitializationResponse;
import org.springframework.vault.support.VaultMount;
import org.springframework.vault.support.VaultToken;
import org.springframework.vault.support.VaultTokenRequest;
import org.springframework.vault.support.VaultTokenResponse;
import org.springframework.vault.support.VaultUnsealStatus;

/**
 * @author Mark Paluch
 */
public class PrepareVault {

	private final VaultOperations vaultOperations;
	private final VaultSysOperations adminOperations;

	public PrepareVault(VaultOperations vaultOperations) {

		this.vaultOperations = vaultOperations;
		this.adminOperations = vaultOperations.opsForSys();
	}

	/**
	 * Initialize Vault and unseal the vault.
	 *
	 * @return the root token.
	 */
	public VaultToken initializeVault() {

		int createKeys = 2;
		int requiredKeys = 2;

		VaultInitializationResponse initialized = vaultOperations.opsForSys()
				.initialize(VaultInitializationRequest.create(createKeys, requiredKeys));

		for (int i = 0; i < requiredKeys; i++) {

			VaultUnsealStatus unsealStatus = vaultOperations.opsForSys()
					.unseal(initialized.getKeys().get(i));

			if (!unsealStatus.isSealed()) {
				break;
			}
		}

		return initialized.getRootToken();
	}

	/**
	 * Create a token for the given {@code tokenId} and {@code policy}.
	 *
	 * @param tokenId
	 * @param policy
	 * @return
	 */
	public VaultToken createToken(String tokenId, String policy) {

		VaultTokenRequest tokenRequest = new VaultTokenRequest();

		tokenRequest.setId(tokenId);
		if (policy != null) {
			tokenRequest.setPolicies(Collections.singletonList(policy));
		}

		VaultTokenResponse vaultTokenResponse = vaultOperations.opsForToken()
				.create(tokenRequest);
		return vaultTokenResponse.getToken();
	}

	/**
	 * Check whether Vault is available (vault created and unsealed).
	 *
	 * @return
	 */
	public boolean isAvailable() {
		return adminOperations.isInitialized() && !adminOperations.health().isSealed();
	}

	/**
	 * Mount an auth backend.
	 *
	 * @param authBackend
	 */
	public void mountAuth(String authBackend) {

		Assert.hasText(authBackend, "AuthBackend must not be empty");

		adminOperations.authMount(authBackend, VaultMount.create(authBackend));
	}

	/**
	 * Check whether a auth-backend is enabled.
	 *
	 * @param authBackend
	 * @return
	 */
	public boolean hasAuth(String authBackend) {

		Assert.hasText(authBackend, "AuthBackend must not be empty");

		return adminOperations.getAuthMounts().containsKey(authBackend + "/");
	}

	/**
	 * Mount an secret backend.
	 *
	 * @param secretBackend
	 */
	public void mountSecret(String secretBackend) {

		Assert.hasText(secretBackend, "SecretBackend must not be empty");

		adminOperations.mount(secretBackend, VaultMount.create(secretBackend));
	}

	/**
	 * Check whether a auth-backend is enabled.
	 *
	 * @param secretBackend
	 * @return
	 */
	public boolean hasSecretBackend(String secretBackend) {

		Assert.hasText(secretBackend, "SecretBackend must not be empty");
		Map<String, VaultMount> mounts = adminOperations.getMounts();
		return mounts.containsKey(secretBackend)
				|| mounts.containsKey(secretBackend + "/");
	}

	public VaultOperations getVaultOperations() {
		return vaultOperations;
	}

	/**
	 * @return Vault version from the health check. Versions beginning from Vault 0.6.2
	 * will expose a version number.
	 */
	public Version getVersion() {

		VaultHealth health = getVaultOperations().opsForSys().health();

		if (StringUtils.hasText(health.getVersion())) {

			String version = health.getVersion();

			// Migration code for Vault 0.6.1
			if (version.startsWith("Vault v")) {
				version = version.substring(7);
			}

			return Version.parse(version);
		}

		return Version.parse("0.0.0");
	}
}

@@ -660,38 +660,8 @@
 l;%0A%0A
-import java.util.Collections;%0A
 impo
@@ -1228,32 +1228,117 @@
 ltTokenRequest;%0A
+import org.springframework.vault.support.VaultTokenRequest.VaultTokenRequestBuilder;%0A
 import org.sprin
@@ -2526,24 +2526,25 @@
  policy) %7B%0A%0A
+%0A
 %09%09VaultToken
@@ -2554,69 +2554,55 @@
 uest
- tokenRequest = new VaultTokenRequest();%0A%0A%09%09tokenRequest.setI
+Builder builder = VaultTokenRequest.builder().i
 d(to
@@ -2613,93 +2613,75 @@
 d);%0A
+%0A
 %09%09if (
-policy != null) %7B%0A%09%09%09tokenRequest.setPolicies(Collections.singletonList
+StringUtils.hasText(policy)) %7B%0A%09%09%09builder.withPolicy
 (policy)
 );%0A%09
@@ -2676,17 +2676,16 @@
 (policy)
-)
 ;%0A%09%09%7D%0A%0A%09
@@ -2771,20 +2771,23 @@
 ate(
-tokenRequest
+builder.build()
 );%0A%09

Use VaultTokenRequestBuilder in Vault preparation.