S3ScriptedPlugin.scala

import sbt._
import Keys._
import com.amazonaws.auth.AWSCredentials
import com.amazonaws.client.builder.AwsClientBuilder.EndpointConfiguration
import com.dimafeng.testcontainers.LocalStackV2Container
import com.amazonaws.services.s3.{AmazonS3, AmazonS3Client}
import com.amazonaws.services.s3.internal.SkipMd5CheckStrategy
import com.amazonaws.services.s3.model.ObjectListing
import org.testcontainers.containers.localstack.LocalStackContainer.Service
import sbt.ScriptedPlugin.autoImport.{scripted, scriptedLaunchOpts}
import sbt.plugins.SbtPlugin
import scala.annotation.tailrec
import scala.collection.JavaConverters._

case object S3ScriptedPlugin extends AutoPlugin {
  override def requires = SbtPlugin
  override def trigger = allRequirements

  val TestBucket = "maven.custom"
  val s3ContainerKey = AttributeKey[LocalStackV2Container]("fm-s3-container")
  val s3ClientKey = AttributeKey[AmazonS3]("fm-s3-client")
  val s3EndpointConfig = AttributeKey[EndpointConfiguration]("fm-s3-container-endpoint-config")
  val s3CredentialsKey = AttributeKey[AWSCredentials]("fm-s3-credentials")

  override lazy val projectSettings = Seq(
    scripted := scripted.dependsOn(setupS3ServerTask).evaluated
  )

  override val globalSettings: Seq[Def.Setting[_]] = Seq(
    Global / onLoad := (Global / onLoad).value andThen startS3Server,
  )

  private def getEndpointConfig(state: State): Option[EndpointConfiguration] = state.get(s3ContainerKey).map{ _.container.getEndpointConfiguration(Service.S3) }
  private def getCredentials(state: State): Option[AWSCredentials] = state.get(s3ContainerKey).map{ _.container.getDefaultCredentialsProvider.getCredentials }
  private def getS3Client(state: State): Option[AmazonS3] = state.get(s3ClientKey)
  private def getServiceEndpoint(state: State): Option[String] = getEndpointConfig(state).map{ _.getServiceEndpoint }
  private def getSigningRegion(state: State): Option[String] = getEndpointConfig(state).map{ _.getSigningRegion }
  private def getAWSAccessKeyId(state: State): Option[String] = getCredentials(state).map{ _.getAWSAccessKeyId }
  private def getAWSSecretKey(state: State): Option[String] = getCredentials(state).map{ _.getAWSSecretKey }

  private def startS3Server(state: State): State = {
    state.get(s3ContainerKey) match {
      case None =>
        state.log.info("[S3ScriptedPlugin] Starting Embedded S3 Server...")
        val container = LocalStackV2Container(services = List(Service.S3))
        container.start()
        state.log.info(s"[S3ScriptedPlugin] Started S3 Server on endpoint: ${container.container.getEndpointConfiguration(Service.S3).getServiceEndpoint}")

        val client = AmazonS3Client.builder()
          .withForceGlobalBucketAccessEnabled(true)
          .withEndpointConfiguration(container.container.getEndpointConfiguration(Service.S3))
          .withCredentials(container.container.getDefaultCredentialsProvider)
          .build()

        val newState = state
          .put(s3ContainerKey, container)
          .put(s3ClientKey, client)
          .put(s3EndpointConfig, container.container.getEndpointConfiguration(Service.S3))
          .put(s3CredentialsKey, container.container.getDefaultCredentialsProvider.getCredentials)
          .addExitHook {
            state
              .get(s3ContainerKey)
              .map { container =>
                state.log.info(s"[S3ScriptedPlugin] Stopping Embedded S3 Server...")
                container.stop()
                state.log.info(s"[S3ScriptedPlugin] Stopping Embedded S3 Server...done")
                state.remove(s3ContainerKey)
              }.getOrElse(state)
          }

        //setupS3Server(newState)

        // This triggers onLoad and onUnload "again", so we must persist a single s3 server instance
        // in-between session reloads, so the settingKey values are correct when the onLoad is ran.
        Project
          .extract(newState)
          .appendWithSession(
            Seq(
              scriptedLaunchOpts ++= Seq(
                getServiceEndpoint(newState).map{ "-Dfm.sbt.s3.endpoint.serviceEndpoint=" + _ }.get,
                getSigningRegion(newState).map{ "-Dfm.sbt.s3.endpoint.signingRegion=" + _ }.get,
                getAWSAccessKeyId(newState).map{ "-Daws.accessKeyId=" + _ }.get,
                getAWSSecretKey(newState).map{ "-Daws.secretKey=" + _ }.get,
                // TODO: The ivy-style publish fails on the ivy.xml MD5 if this isn't set, I think its
                //       specific to localstack and not production Amazon S3
                "-D"+SkipMd5CheckStrategy.DISABLE_PUT_OBJECT_MD5_VALIDATION_PROPERTY+"=true"
              )
            ),
            newState
          )
      case Some(container) =>
        state.log.info(s"[S3ScriptedPlugin] Reusing S3 Server ${container.container.getEndpointConfiguration(Service.S3).getServiceEndpoint}")
        state
    }
  }

  private lazy val setupS3ServerTask: Def.Initialize[Task[Unit]] = Def.task {
    setupS3Server(state.value)
  }

  @tailrec private def deleteS3BucketObjects(client: AmazonS3, bucketName: String, objectListing: ObjectListing): Unit = {
    objectListing.getObjectSummaries.iterator().asScala.foreach { obj =>
      client.deleteObject(bucketName, obj.getKey)
    }
    if (objectListing.isTruncated) deleteS3BucketObjects(client, bucketName, client.listNextBatchOfObjects(objectListing))
  }

  private def setupS3Server(state: State): Unit = {
    getS3Client(state).foreach { client =>
      state.log.info(s"[S3ScriptedPlugin] Recreating s3 bucket: '$TestBucket'")
      client.listBuckets().asScala.foreach{ bucket =>
        val bucketName = bucket.getName
        deleteS3BucketObjects(client, bucketName, client.listObjects(bucketName))
        client.deleteBucket(bucketName)
      }
      client.createBucket(TestBucket)
    }
  }
//  //onUnload in Global := (onUnload in Global).value andThen finish
//  private def finish(state: State): State = {
//    state
//      .get(s3ContainerKey)
//      .map { container =>
//        state.log.info(s"[S3ScriptedPlugin] Stopping Embedded S3 Server...")
//        container.stop()
//        state.log.info(s"[S3ScriptedPlugin] Stopping Embedded S3 Server...done")
//        state.remove(s3ContainerKey)
//      }
//      .getOrElse(state)
//  }
}