D-Net

package eu.dnetlib.msro.workers.aggregation.collect.plugins.ftp;

import java.io.IOException;

import java.io.OutputStream;

import java.net.MalformedURLException;

import java.net.URL;

import java.util.Iterator;

import java.util.LinkedList;

import java.util.Queue;

import java.util.Set;

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

import org.apache.commons.logging.Log;

import org.apache.commons.logging.LogFactory;

import org.apache.commons.net.ftp.FTPClient;

import org.apache.commons.net.ftp.FTPFile;

import org.apache.commons.net.ftp.FTPReply;

import eu.dnetlib.msro.exceptions.CollectorServiceRuntimeException;

/**

 * @author Author: Andrea Mannocci

 */

public class FtpIterator implements Iterator<String> {

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

	private static final int MAX_RETRIES = 5;

	private static final int DEFAULT_TIMEOUT = 30000;

	private static final long BACKOFF_MILLIS = 10000;

	private FTPClient ftpClient;

	private String ftpServerAddress;

	private String remoteFtpBasePath;

	private String username;

	private String password;

	private boolean isRecursive;

	private Set<String> extensionsSet;

	private Queue<String> queue;

	public FtpIterator(final String baseUrl, final String username, final String password, final boolean isRecursive,

			final Set<String> extensionsSet) {

		this.username = username;

		this.password = password;

		this.isRecursive = isRecursive;

		this.extensionsSet = extensionsSet;

		try {

			final URL server = new URL(baseUrl);

			ftpServerAddress = server.getHost();

			remoteFtpBasePath = server.getPath();

		} catch (final MalformedURLException e1) {

			throw new CollectorServiceRuntimeException("Malformed URL exception " + baseUrl);

		}

		connectToFtpServer();

		initializeQueue();

	}

	private void connectToFtpServer() {

		ftpClient = new FTPClient();

		ftpClient.setDefaultTimeout(DEFAULT_TIMEOUT);

		ftpClient.setDataTimeout(DEFAULT_TIMEOUT);

		ftpClient.setConnectTimeout(DEFAULT_TIMEOUT);

		try {

			ftpClient.connect(ftpServerAddress);

			// try to login

			if (!ftpClient.login(username, password)) {

				ftpClient.logout();

				throw new CollectorServiceRuntimeException("Unable to login to FTP server " + ftpServerAddress);

			}

			final int reply = ftpClient.getReplyCode();

			if (!FTPReply.isPositiveCompletion(reply)) {

				ftpClient.disconnect();

				throw new CollectorServiceRuntimeException("Unable to connect to FTP server " + ftpServerAddress);

			}

			ftpClient.enterLocalPassiveMode();

			log.info("Connected to FTP server " + ftpServerAddress);

			log.info(String.format("FTP collecting from %s with recursion = %s", remoteFtpBasePath, isRecursive));

		} catch (final IOException e) {

			throw new CollectorServiceRuntimeException("Unable to connect to FTP server " + ftpServerAddress);

		}

	}

	private void disconnectFromFtpServer() {

		try {

			if (ftpClient.isConnected()) {

				ftpClient.logout();

				ftpClient.disconnect();

			}

		} catch (final IOException e) {

			log.error("Failed to logout & disconnect from the FTP server", e);

		}

	}

	private void initializeQueue() {

		queue = new LinkedList<String>();

		listDirectoryRecursive(remoteFtpBasePath, "");

	}

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

		String dirToList = parentDir;

		if (!currentDir.equals("")) {

			dirToList += "/" + currentDir;

		}

		FTPFile[] subFiles;

		try {

			subFiles = ftpClient.listFiles(dirToList);

			if ((subFiles != null) && (subFiles.length > 0)) {

				for (final FTPFile aFile : subFiles) {

					final String currentFileName = aFile.getName();

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

						// skip parent directory and directory itself

						continue;

					}

					if (aFile.isDirectory()) {

						if (isRecursive) {

							listDirectoryRecursive(dirToList, currentFileName);

						}

					} else {

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

						for (final String ext : extensionsSet) {

							if (currentFileName.endsWith(ext)) {

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

							}

						}

					}

				}

			}

		} catch (final IOException e) {

			throw new CollectorServiceRuntimeException("Unable to list FTP remote folder", e);

		}

	}

	@Override

	public boolean hasNext() {

		if (queue.isEmpty()) {

			disconnectFromFtpServer();

			return false;

		} else {

			return true;

		}

	}

	@Override

	public String next() {

		final String nextRemotePath = queue.remove();

		int nRepeat = 0;

		while (nRepeat < MAX_RETRIES) {

			try {

				final OutputStream baos = new ByteArrayOutputStream();

				if (!ftpClient.isConnected()) {

					connectToFtpServer();

				}

				ftpClient.retrieveFile(nextRemotePath, baos);

				log.debug(String.format("Collected file from FTP: %s%s", ftpServerAddress, nextRemotePath));

				return baos.toString();

			} catch (final IOException e) {

				nRepeat++;

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

						nRepeat));

				disconnectFromFtpServer();

				try {

					Thread.sleep(BACKOFF_MILLIS);

				} catch (final InterruptedException e1) {

					log.error(e1);

				}

			}

		}

		throw new CollectorServiceRuntimeException(

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

	}

	@Override

	public void remove() {

		throw new UnsupportedOperationException();

	}

}