D-Net

package eu.dnetlib.data.collector.plugins.sftp;

import java.io.OutputStream;

import java.net.URI;

import java.net.URISyntaxException;

import java.time.Instant;

import java.time.LocalDate;

import java.time.LocalDateTime;

import java.time.ZoneId;

import java.time.format.DateTimeFormatter;

import java.time.format.DateTimeParseException;

import java.util.LinkedList;

import java.util.Queue;

import java.util.Set;

import java.util.Vector;

import com.jcraft.jsch.*;

import eu.dnetlib.data.collector.ThreadSafeIterator;

import eu.dnetlib.rmi.data.CollectorServiceRuntimeException;

import org.apache.commons.io.output.ByteArrayOutputStream;

import org.apache.commons.lang3.StringUtils;

import org.apache.commons.logging.Log;

import org.apache.commons.logging.LogFactory;

/**

 * Created by andrea on 11/01/16.

 */

public class SftpIterator extends ThreadSafeIterator {

	private static final Log log = LogFactory.getLog(SftpIterator.class);

	private static final int MAX_RETRIES = 2;

	private static final int DEFAULT_TIMEOUT = 30000;

	private static final long BACKOFF_MILLIS = 10000;

	//params for simple authentication mode

	private String username;

	private String password;

	//params for pubkey authentication mode

	private String prvKeyFile;

	private String passPhrase;

	private String knownHostsFile;

	private String baseUrl;

	private String sftpURIScheme;

	private String sftpServerAddress;

	private String remoteSftpBasePath;

	private boolean isRecursive;

	private Set<String> extensionsSet;

	private boolean incremental;

	private Session sftpSession;

	private ChannelSftp sftpChannel;

	private Queue<String> queue;

	private LocalDateTime fromDate = null;

	private DateTimeFormatter simpleDateTimeFormatter = DateTimeFormatter.ofPattern("yyyy-MM-dd");

	private DateTimeFormatter fullDateTimeFormatter = DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss'Z'");

	private void init(final String baseUrl, final boolean isRecursive, final Set<String> extensionsSet, final String fromDate) {

		this.baseUrl = baseUrl;

		this.isRecursive = isRecursive;

		this.extensionsSet = extensionsSet;

		this.incremental = StringUtils.isNotBlank(fromDate);

		if (incremental) {

			try {

				this.fromDate = LocalDateTime.parse(fromDate, fullDateTimeFormatter);

			}catch(DateTimeParseException dtpe){

				//try with the simple formatter

				this.fromDate = LocalDateTime.from(LocalDate.parse(fromDate, simpleDateTimeFormatter).atStartOfDay());

			}

			log.debug("fromDate string: " + fromDate + " -- parsed: " + this.fromDate.toString());

		}

		try {

			URI sftpServer = new URI(this.baseUrl);

			this.sftpURIScheme = sftpServer.getScheme();

			this.sftpServerAddress = sftpServer.getHost();

			this.remoteSftpBasePath = sftpServer.getPath();

		} catch (URISyntaxException e) {

			throw new CollectorServiceRuntimeException("Bad syntax in the URL " + baseUrl);

		}

	}

	public SftpIterator(String baseUrl, String username, String password, boolean isRecursive, Set<String> extensionsSet, String fromDate) {

		init(baseUrl, isRecursive, extensionsSet, fromDate);

		this.username = username;

		this.password = password;

		connectToSftpServerSimpleAuth();

		initializeQueue();

	}

	public SftpIterator(final String baseUrl, final String username, final String prvKeyFilePath, final String passPhrase, final String knownHostsFile, final boolean isRecursive, final Set<String> extensionsSet, final String fromDate) {

		init(baseUrl, isRecursive, extensionsSet, fromDate);

		this.username = username;

		this.passPhrase = passPhrase;

		this.prvKeyFile = prvKeyFilePath;

		this.knownHostsFile = knownHostsFile;

		connectToSftpServerPubKeyAuth();

		initializeQueue();

	}

	private void connectToSftpServerPubKeyAuth() {

		JSch jsch = new JSch();

		log.info("Connecting to "+sftpServerAddress+" with PubKey authentication");

		log.info("Username "+ username);

		log.info("Private key path: "+prvKeyFile);

		if(StringUtils.isNotBlank(passPhrase)){

			log.info("with Pass phrase");

		}

		log.info("Known host file path: "+knownHostsFile);

		try {

			jsch.setKnownHosts(this.knownHostsFile);

			jsch.addIdentity(this.prvKeyFile, this.passPhrase);

			sftpSession = jsch.getSession(username, sftpServerAddress);

			sftpSession.connect();

			openChannelOnBasePath();

		} catch (JSchException e) {

			throw new CollectorServiceRuntimeException("Unable to create a session on remote SFTP server via Public key authentication.", e);

		}

	}

	private void connectToSftpServerSimpleAuth() {

		JSch jsch = new JSch();

		try {

			JSch.setConfig("StrictHostKeyChecking", "no");

			sftpSession = jsch.getSession(username, sftpServerAddress);

			sftpSession.setPassword(password);

			sftpSession.connect();

			openChannelOnBasePath();

		} catch (JSchException e) {

			throw new CollectorServiceRuntimeException("Unable to create a session on remote SFTP server via simple authentication.", e);

		}

	}

	private void openChannelOnBasePath() {

		String fullPath = "";

		try {

			Channel channel = sftpSession.openChannel(sftpURIScheme);

			channel.connect();

			sftpChannel = (ChannelSftp) channel;

			String pwd = sftpChannel.pwd();

			log.debug("PWD from server: " + pwd);

			fullPath = pwd + remoteSftpBasePath;

			sftpChannel.cd(fullPath);

			log.debug("PWD from server 2 after 'cd " + fullPath + "' : " + sftpChannel.pwd());

			log.info("Connected to SFTP server " + sftpServerAddress);

		} catch (JSchException e) {

			throw new CollectorServiceRuntimeException("Unable to open/connect SFTP channel.", e);

		} catch (SftpException e) {

			throw new CollectorServiceRuntimeException("Unable to access the remote path "+fullPath+" on the SFTP server.", e);

		}

	}

	private void disconnectFromSftpServer() {

		sftpChannel.exit();

		sftpSession.disconnect();

	}

	private void initializeQueue() {

		queue = new LinkedList<>();

		log.info(String.format("SFTP collector plugin collecting from %s with recursion = %s, incremental = %s with fromDate=%s", remoteSftpBasePath,

				isRecursive,

				incremental, fromDate));

		listDirectoryRecursive(".", "");

	}

	private void listDirectoryRecursive(final String parentDir, final String currentDir) {

		String dirToList = parentDir;

		if (StringUtils.isNotBlank(currentDir)) {

			dirToList += "/" + currentDir;

		}

		log.debug("PARENT DIR: " + parentDir);

		log.debug("DIR TO LIST: " + dirToList);

		try {

			Vector<ChannelSftp.LsEntry> ls = sftpChannel.ls(dirToList);

			for (ChannelSftp.LsEntry entry : ls) {

				String currentFileName = entry.getFilename();

				if (currentFileName.equals(".") || currentFileName.equals("..")) {

					// skip parent directory and directory itself

					continue;

				}

				SftpATTRS attrs = entry.getAttrs();

				if (attrs.isDir()) {

					if (isRecursive) {

						listDirectoryRecursive(dirToList, currentFileName);

					}

				} else {

					// test the file for extensions compliance and, just in case, add it to the list.

					for (String ext : extensionsSet) {

						if (currentFileName.endsWith(ext)) {

							//test if the file has been changed after the last collection date:

							if (incremental) {

								int mTime = attrs.getMTime();

								//int times are values reduced by the milliseconds, hence we multiply per 1000L

								Instant instant = Instant.ofEpochMilli(mTime * 1000L);

								LocalDateTime dt = LocalDateTime.ofInstant(instant, ZoneId.systemDefault());

								if (dt.isAfter(fromDate)) {

									queue.add(dirToList+"/"+currentFileName);

									log.debug(dirToList+"/"+currentFileName + " has changed and must be re-collected");

								} else {

									if (log.isDebugEnabled()) {

										log.debug(dirToList+"/"+currentFileName + " has not changed since last collection");

									}

								}

							} else {

								//if it is not incremental, just add it to the queue

								queue.add(dirToList+"/"+currentFileName);

							}

						}

					}

				}

			}

		} catch (SftpException e) {

			throw new CollectorServiceRuntimeException("Cannot list the sftp remote directory", e);

		}

	}

	@Override

	public boolean doHasNext() {

		return !queue.isEmpty();

	}

	@Override

	public String doNext() {

		if(queue.isEmpty())

			throw new CollectorServiceRuntimeException("Unexpected empty queue in next()");

		String nextRemotePath = queue.remove();

		int nRepeat = 0;

		String fullPathFile = "";

		while (nRepeat < MAX_RETRIES) {

			try {

				OutputStream baos = new ByteArrayOutputStream();

				sftpChannel.get(nextRemotePath, baos);

				if (log.isDebugEnabled()) {

					fullPathFile  = sftpChannel.pwd() + "/" + nextRemotePath;

					log.debug(String.format("Collected file from SFTP: %s%s", sftpServerAddress, fullPathFile));

				}

				if (queue.isEmpty()) {

					disconnectFromSftpServer();

				}

				return baos.toString();

			} catch (SftpException e) {

				nRepeat++;

				log.warn(String.format("An error occurred [%s] for %s%s, retrying.. [retried %s time(s)]", e.getMessage(), sftpServerAddress, fullPathFile,

						nRepeat));

				e.printStackTrace();

				try {

					Thread.sleep(BACKOFF_MILLIS);

				} catch (InterruptedException e1) {

					log.error(e1);

				}

			}

		}

		throw new CollectorServiceRuntimeException(

				String.format("Impossible to retrieve FTP file %s after %s retries. Aborting FTP collection.", fullPathFile, nRepeat));

	}

	@Override

	public void remove() {

		throw new UnsupportedOperationException();

	}

}