package io.fathom.cloud.compute.scheduler;
import io.fathom.cloud.sftp.RemoteFile;
import io.fathom.cloud.sftp.RemoteTempFile;
import io.fathom.cloud.ssh.SftpChannel;
import io.fathom.cloud.ssh.SshConfig;
import io.fathom.cloud.ssh.SftpChannel.WriteMode;
import java.io.File;
import java.io.IOException;
import java.io.OutputStream;
import com.google.common.base.Charsets;
import com.google.common.io.ByteSource;
public class WriteFile {
private final SshConfig sshConfig;
private File dest;
private ByteSource source;
private Integer chmod;
private Integer chownUser;
private Integer chownGroup;
private boolean sudo;
private boolean atomic;
private RemoteFile atomicTempPath;
WriteFile(SshConfig sshConfig) {
this.sshConfig = sshConfig;
}
public static WriteFile with(SshConfig sshConfig) {
return new WriteFile(sshConfig);
}
public WriteFile to(File dest) {
this.dest = dest;
return this;
}
public WriteFile from(byte[] data) {
this.source = ByteSource.wrap(data);
return this;
}
public WriteFile from(String s) {
return from(s.getBytes(Charsets.UTF_8));
}
public WriteFile chmod(int mode) {
this.chmod = mode;
return this;
}
// chown and chgrp require CAP_CHOWN / root
// It's a huge backdoor
@Deprecated
public WriteFile chown(int user, int group) {
this.chownUser = user;
this.chownGroup = group;
return this;
}
public WriteFile withSudo() {
this.sudo = true;
return this;
}
public WriteFile atomic(RemoteFile tempPath) {
this.atomicTempPath = tempPath;
this.atomic = true;
return this;
}
public void run() throws IOException {
if (dest == null || source == null) {
throw new IllegalStateException();
}
try (SftpChannel sftp = sshConfig.getSftpChannel()) {
RemoteTempFile tempFile = null;
File remoteDest = dest;
try {
if (sudo || atomic) {
if (atomicTempPath != null) {
tempFile = RemoteTempFile.create(sftp, atomicTempPath);
} else {
tempFile = RemoteTempFile.create(sftp, new RemoteFile(new File("/tmp")));
}
remoteDest = tempFile.getSshPath();
}
try (OutputStream os = sftp.writeFile(remoteDest, WriteMode.Overwrite)) {
source.copyTo(os);
}
if (chmod != null) {
sftp.chmod(remoteDest, chmod);
}
// chown and chgrp require CAP_CHOWN / root
if (chownGroup != null) {
sftp.chgrp(remoteDest, chownGroup);
}
if (chownUser != null) {
sftp.chown(remoteDest, chownUser);
}
if (tempFile != null) {
if (atomic) {
if (sudo) {
throw new UnsupportedOperationException("Atomic sudo move not yet implemented");
}
tempFile.renameTo(new RemoteFile(dest));
} else {
ShellCommand command = ShellCommand.create("/bin/cp");
command.arg(remoteDest);
command.arg(dest);
if (sudo) {
command.useSudo();
}
SshCommand sshCommand = command.withSsh(sshConfig);
sshCommand.run();
}
}
} finally {
if (tempFile != null) {
tempFile.close();
}
}
}
}
}