OpenSearch/plugin/ml-cpp-snapshot/build.gradle

import com.amazonaws.AmazonServiceException
import com.amazonaws.ClientConfiguration
import com.amazonaws.auth.AWSCredentials
import com.amazonaws.auth.BasicAWSCredentials
import com.amazonaws.services.s3.AmazonS3Client
import com.amazonaws.services.s3.model.S3Object
import com.amazonaws.services.s3.model.ObjectMetadata
import com.bettercloud.vault.Vault
import com.bettercloud.vault.VaultConfig
import com.bettercloud.vault.response.LogicalResponse
import java.nio.file.Files
import java.nio.file.attribute.PosixFilePermission
import java.nio.file.attribute.PosixFilePermissions
import org.elasticsearch.gradle.VersionProperties

apply plugin: 'distribution'

buildscript {
  repositories {
    mavenCentral()
  }
  dependencies {
    classpath group: 'com.bettercloud', name: 'vault-java-driver', version:"1.1.0"
    classpath 'com.amazonaws:aws-java-sdk-s3:1.10.33'
    if (JavaVersion.current() > JavaVersion.VERSION_1_8) {
      classpath 'com.sun.xml.bind:jaxb-impl:2.2.3-1' // pulled in as external dependency to work on java 9
    }
  }
}

ext.version = VersionProperties.elasticsearch

// This project pulls a snapshot version of the ML cpp artifacts and sets that as the artifact
// for this project so it can be used with dependency substitution. We do not use gradle's
// handling of S3 as a maven repo due to the dynamically generated creds being slow to propagate,
// necessitating retries.

void checkJavaVersion() {
  /**
   * The Elastic Secrets vault is served via HTTPS with a Let's Encrypt certificate. The root certificates that cross-signed the Let's
   * Encrypt certificates were not trusted by the JDK until 8u101. Therefore, we enforce that the JDK is at least 8u101 here.
   */
  final String javaVersion = System.getProperty('java.version')
  final String javaVendor = System.getProperty('java.vendor')
  def matcher = javaVersion =~ /1\.8\.0(?:_(\d+))?/
  boolean matches = matcher.matches()
  if (matches) {
    final int update
    if (matcher.group(1) == null) {
      update = 0
    } else {
      update = matcher.group(1).toInteger()
    }
    if (update < 101) {
      throw new GradleException("JDK ${javaVendor} ${javaVersion} does not have necessary root certificates " +
              "(https://bugs.openjdk.java.net/browse/JDK-8154757), update your JDK to at least JDK 8u101+")
    }
  }
}

void setupVaultAuthMethod() {
  String VAULT_BASE_URL = 'https://secrets.elastic.co:8200'
  String VAULT_ROLE_ID = "8e90dd88-5a8e-9c12-0da9-5439f293ff97"
  String VAULT_SECRET_ID = System.env.VAULT_SECRET_ID
  // get an authentication token with vault
  String homePath = System.properties['user.home']
  File githubToken = file("${homePath}/.elastic/github.token")
  String vaultAuthBody = null
  URL vaultUrl = null
  if (githubToken.exists()) {
    try {
      Set<PosixFilePermission> perms = Files.getPosixFilePermissions(githubToken.toPath())
      if (perms.equals(PosixFilePermissions.fromString("rw-------")) == false) {
        throw new GradleException('github.token must have 600 permissions')
      }
    } catch (UnsupportedOperationException e) {
      // Assume this isn't a POSIX file system
    }
    vaultUrl = new URL(VAULT_BASE_URL + '/v1/auth/github/login')
    vaultAuthBody = "{\"token\": \"${githubToken.getText('UTF-8').trim()}\"}"
  } else if (VAULT_SECRET_ID != null) {
    vaultUrl = new URL(VAULT_BASE_URL + '/v1/auth/approle/login')
    vaultAuthBody = "{\"role_id\": \"${VAULT_ROLE_ID}\", \"secret_id\": \"${VAULT_SECRET_ID}\"}"
  } else {
    throw new GradleException('Missing ~/.elastic/github.token file or VAULT_SECRET_ID environment variable, needed to authenticate with vault for secrets')
  }
  project.ext.vaultAuthBody = vaultAuthBody
  project.ext.vaultUrl = vaultUrl
}

void getZip(File snapshotZip) {
  HttpURLConnection vaultConn = (HttpURLConnection) vaultUrl.openConnection()
  vaultConn.setRequestProperty('Content-Type', 'application/json')
  vaultConn.setRequestMethod('PUT')
  vaultConn.setDoOutput(true)
  vaultConn.outputStream.withWriter('UTF-8') { writer ->
    writer.write(vaultAuthBody)
  }
  vaultConn.connect()
  Object authResponse = new groovy.json.JsonSlurper().parseText(vaultConn.content.text)
  VaultConfig config = new VaultConfig('https://secrets.elastic.co:8200', authResponse.auth.client_token)
  Vault vault = new Vault(config)
  LogicalResponse secret = vault.logical().read("aws-dev/creds/prelertartifacts")
  final AWSCredentials creds = new BasicAWSCredentials(secret.data.get('access_key'), secret.data.get('secret_key'))

  // the keys may take a while to propagate, so wait up to 60 seconds retrying
  final AmazonS3Client client = new AmazonS3Client(creds)
  final String key = "maven/org/elasticsearch/ml/ml-cpp/${version}/ml-cpp-${version}.zip"
  int retries = 120
  while (retries > 0) {
    try {
      File snapshotMd5 = new File(snapshotZip.toString() + '.md5')
      // do a HEAD first to check the zip hash against the local file
      ObjectMetadata metadata = client.getObjectMetadata('prelert-artifacts', key)
      String remoteMd5 = metadata.getETag()
      if (snapshotZip.exists()) {
        // do a HEAD first to check the zip hash against the local file
        String localMd5 = snapshotMd5.getText('UTF-8')
        if (remoteMd5.equals(localMd5)) {
          logger.info('Using cached ML snapshot')
          return
        }
      }
      S3Object zip = client.getObject('prelert-artifacts', key)
      InputStream zipStream = zip.getObjectContent()
      try {
        project.delete(snapshotZip, snapshotZip)
        Files.copy(zipStream, snapshotZip.toPath())
      } finally {
        zipStream.close()
      }
      snapshotMd5.setText(remoteMd5, 'UTF-8')
      return
    } catch (AmazonServiceException e) {
      if (e.getStatusCode() != 403) {
        throw new GradleException('Error while trying to get ml-cpp snapshot: ' + e.getMessage(), e)
      }
      sleep(500)
      retries--
    }
  }
  throw new GradleException('Could not access ml-cpp artifacts. Timed out after 60 seconds')
}

File snapshotZip = new File(projectDir, ".cache/ml-cpp-${version}.zip")
task downloadMachineLearningSnapshot {
  onlyIf {
    // skip if machine-learning-cpp is being built locally
    findProject(':machine-learning-cpp') == null &&
      // skip for offline builds - just rely on the artifact already having been downloaded before here
      project.gradle.startParameter.isOffline() == false
  }
  doFirst {
    snapshotZip.parentFile.mkdirs()
    getZip(snapshotZip)
  }
}

gradle.taskGraph.whenReady { taskGraph ->
  // skip if machine-learning-cpp is being built locally and also for offline builds
  if (findProject(':machine-learning-cpp') == null && project.gradle.startParameter.isOffline() == false) {
    // do validation of token/java version up front, don't wait for the task to run
    checkJavaVersion()
    setupVaultAuthMethod()
  }
}

artifacts {
  'default' file: snapshotZip, name: 'ml-cpp', type: 'zip', builtBy: downloadMachineLearningSnapshot
}