package com.hubspot.singularity.executor.task;

import com.hubspot.deploy.S3Artifact;
import com.hubspot.deploy.S3ArtifactSignature;
import com.hubspot.singularity.executor.config.SingularityExecutorConfiguration;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import org.slf4j.Logger;

public class SingularityExecutorArtifactVerifier {
  private final Logger log;
  private final SingularityExecutorConfiguration executorConfiguration;
  private final SingularityExecutorTaskDefinition taskDefinition;

  public SingularityExecutorArtifactVerifier(
    SingularityExecutorTaskDefinition taskDefinition,
    Logger log,
    SingularityExecutorConfiguration executorConfiguration
  ) {
    this.log = log;
    this.executorConfiguration = executorConfiguration;
    this.taskDefinition = taskDefinition;
  }

  public void checkSignatures(
    SingularityExecutorTask task,
    List<S3Artifact> s3Artifacts,
    List<S3ArtifactSignature> s3ArtifactsWithSignatures
  ) {
    if (s3ArtifactsWithSignatures.isEmpty()) {
      log.info(
        "No files containing artifact signatures specified, skipping verification."
      );
      return;
    }

    for (S3ArtifactSignature s3ArtifactSignature : s3ArtifactsWithSignatures) {
      Optional<S3Artifact> maybeMatchingForSignature = s3Artifacts
        .stream()
        .filter(s -> s3ArtifactSignature.getArtifactFilename().equals(s.getFilename()))
        .findFirst();
      if (maybeMatchingForSignature.isPresent()) {
        checkArtifactSignature(
          task,
          maybeMatchingForSignature.get(),
          s3ArtifactSignature
        );
      } else {
        log.warn("No matching artifact found for signature {}", s3ArtifactSignature);
        if (executorConfiguration.isFailOnSignatureWithNoMatchingArtifact()) {
          throw new RuntimeException(
            String.format(
              "No matching artifact found for signature %s",
              s3ArtifactSignature.getFilename()
            )
          );
        }
      }
    }
  }

  private void checkArtifactSignature(
    SingularityExecutorTask task,
    S3Artifact s3Artifact,
    S3ArtifactSignature s3ArtifactSignature
  ) {
    final Path artifactPath = task.getArtifactPath(
      s3Artifact,
      task.getTaskDefinition().getTaskDirectoryPath()
    );
    final Path artifactSignaturePath = task.getArtifactPath(
      s3ArtifactSignature,
      task.getTaskDefinition().getTaskDirectoryPath()
    );

    if (!Files.exists(artifactPath)) {
      log.warn(
        "Artifact {} not found for signature {}",
        artifactPath,
        s3ArtifactSignature
      );
      return;
    }

    final List<String> verifyCommand = new ArrayList<>(
      executorConfiguration.getArtifactSignatureVerificationCommand().size()
    );

    for (String arg : executorConfiguration.getArtifactSignatureVerificationCommand()) {
      verifyCommand.add(
        arg
          .replace("{artifactPath}", artifactPath.toString())
          .replace("{artifactSignaturePath}", artifactSignaturePath.toString())
      );
    }

    try {
      final ProcessBuilder processBuilder = new ProcessBuilder(verifyCommand);

      processBuilder.directory(taskDefinition.getTaskDirectoryPath().toFile());

      processBuilder.redirectError(taskDefinition.getSignatureVerifyOutPath().toFile());
      processBuilder.redirectOutput(taskDefinition.getSignatureVerifyOutPath().toFile());

      final Process p = processBuilder.start();

      p.waitFor(); // TODO: add some sort of timeout?

      if (p.exitValue() != 0) {
        log.error(
          "Failed to validate signature in file {} for artifact file {}",
          s3ArtifactSignature.getFilename(),
          s3ArtifactSignature.getArtifactFilename()
        );

        if (executorConfiguration.isFailTaskOnInvalidArtifactSignature()) {
          throw new RuntimeException(
            String.format("Failed to validate signature for artifact %s", artifactPath)
          );
        }
      } else {
        log.info(
          "Signature in {} for artifact {} is valid!",
          s3ArtifactSignature.getFilename(),
          s3ArtifactSignature.getArtifactFilename()
        );
      }
    } catch (InterruptedException | IOException e) {
      throw new RuntimeException(e);
    }
  }
}

@@ -2298,32 +2298,39 @@
 ifactPath = task
+%0A      
 .getArtifactPath
@@ -2322,39 +2322,32 @@
 getArtifactPath(
-%0A      
 s3Artifact,%0A    
@@ -2341,22 +2341,16 @@
 rtifact,
-%0A     
  task.ge
@@ -2389,21 +2389,57 @@
 ryPath()
+)
 %0A    
+  .resolve(s3Artifact.getFilename()
 );%0A    f
@@ -2472,24 +2472,31 @@
 ePath = task
+%0A      
 .getArtifact
@@ -2501,16 +2501,18 @@
 ctPath(%0A
+  
       s3
@@ -2530,16 +2530,18 @@
 nature,%0A
+  
       ta
@@ -2586,27 +2586,78 @@
 yPath()%0A    
+ 
+ )%0A      .resolve(s3ArtifactSignature.getFilename()
 );
-%0A
 %0A    if (!Fi

resolve to filename