ProductResender.java

package gov.usgs.earthquake.distribution;

import java.io.BufferedReader;
import java.io.File;
import java.io.InputStreamReader;
import java.security.PrivateKey;
import java.util.Iterator;
import java.util.Map;
import java.util.logging.Level;
import java.util.logging.Logger;

import gov.usgs.earthquake.aws.AwsProductSender;
import gov.usgs.earthquake.product.Product;
import gov.usgs.earthquake.product.io.IOUtil;
import gov.usgs.earthquake.product.io.ObjectProductHandler;
import gov.usgs.earthquake.product.io.ProductSource;
import gov.usgs.util.CryptoUtils;
import gov.usgs.util.FileUtils;

/**
 * A utility class to (re)send an existing product to pdl hubs.
 *
 * Mainly used when one server has not received a product, in order to
 * redistribute the product.
 */
public class ProductResender {

    private static final Logger LOGGER = Logger.getLogger(ProductResender.class
            .getName());

    /** Servers arguments */
    public static final String SERVERS_ARGUMENT = "--servers=";
    /** Batch arguments */
    public static final String BATCH_ARGUMENT = "--batch";
    /** Private Key Argument */
    public static final String PRIVATE_KEY_ARGUMENT = "--privateKey=";

    /**
     * Command Line Interface to ProductResender.
     *
     * @param args CLI arguments
     * @throws Exception if error occurs
     */
    public static void main(final String[] args) throws Exception {
        // disable tracker
        ProductTracker.setTrackerEnabled(false);

        File inFile = null;
        String inFormat = null;
        String servers = null;
        boolean binaryFormat = false;
        boolean enableDeflate = true;
        boolean batchMode = false;
        PrivateKey privateKey = null;

        for (String arg : args) {
            if (arg.startsWith(IOUtil.INFILE_ARGUMENT)) {
                inFile = new File(arg.replace(IOUtil.INFILE_ARGUMENT, ""));
            } else if (arg.startsWith(IOUtil.INFORMAT_ARGUMENT)) {
                inFormat = arg.replace(IOUtil.INFORMAT_ARGUMENT, "");
            } else if (arg.startsWith(SERVERS_ARGUMENT)) {
                servers = arg.replace(SERVERS_ARGUMENT, "");
            } else if (arg.equals(CLIProductBuilder.BINARY_FORMAT_ARGUMENT)) {
                binaryFormat = true;
            } else if (arg.equals(CLIProductBuilder.DISABLE_DEFLATE)) {
                enableDeflate = false;
            } else if (arg.equals(BATCH_ARGUMENT)) {
                batchMode = true;
            } else if (arg.startsWith(PRIVATE_KEY_ARGUMENT)) {
                privateKey = CryptoUtils.readOpenSSHPrivateKey(
                        FileUtils.readFile(new File(arg.replace(PRIVATE_KEY_ARGUMENT, ""))),
                        null);
                if (privateKey == null) {
                    LOGGER.warning("Unable to parse private key " + arg);
                    System.exit(1);
                }
            }
        }

        // read product
        Product product = null;
        if (!batchMode) {
            ProductSource source = IOUtil.getProductSource(inFormat, inFile);
            if (source != null) {
                product = ObjectProductHandler.getProduct(source);
            }
        }

        ProductBuilder builder = new ProductBuilder();
        builder.getProductSenders().addAll(
                CLIProductBuilder.parseServers(servers, 15000, binaryFormat,
                        enableDeflate));
        if (privateKey != null) {
            // resign products
            for (ProductSender sender : builder.getProductSenders()) {
                if (sender instanceof AwsProductSender) {
                    AwsProductSender awsSender = (AwsProductSender) sender;
                    awsSender.setPrivateKey(privateKey);
                    awsSender.setSignProducts(true);;
                }
            }
        }

        if ((!batchMode && product == null) || builder.getProductSenders().size() == 0) {
            System.err.println("Usage: ProductResender --servers=SERVERLIST"
                    + " --informat=(zip|directory|xml) --infile=FILE"
                    + " [--binaryFormat] [--disableDeflate] [--batch]");
            System.err.println("When using batch mode (--batch), the --infile argument is ignored.");
            System.err.println("Files to send are read one per line from stdin.");
            System.exit(CLIProductBuilder.EXIT_INVALID_ARGUMENTS);
        }

        builder.startup();

        if (!batchMode) {
            sendProduct(builder, product, batchMode);
        } else {
            // send batch
            BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
            String line = null;
            while ((line = br.readLine()) != null) {
                inFile = new File(line);
                ProductSource source = IOUtil.getProductSource(inFormat, inFile);
                if (source != null) {
                    product = ObjectProductHandler.getProduct(source);
                } else {
                    System.err.println("ERROR: unable to load product from '" +
                            inFile.getCanonicalPath() + "'");
                    continue;
                }
                sendProduct(builder, product, batchMode);
            }
        }

        // normal exit
        builder.shutdown();
        System.exit(0);
    }

    /**
     * Sends product to builder
     * @param builder ProductBuilder
     * @param product Product
     * @param batchMode bool
     * @throws Exception if error occurs
     */
    protected static void sendProduct(final ProductBuilder builder,
            final Product product, final boolean batchMode) throws Exception {
        // extracted from CLIProductBuilder

        // send the product
        Map<ProductSender, Exception> sendExceptions = builder
                .sendProduct(product);

        // handle any send exceptions
        if (sendExceptions.size() != 0) {
            Iterator<ProductSender> senders = sendExceptions.keySet()
                    .iterator();
            // log the exceptions
            while (senders.hasNext()) {
                ProductSender sender = senders.next();
                if (sender instanceof SocketProductSender) {
                    // put more specific information about socket senders
                    SocketProductSender socketSender = (SocketProductSender) sender;
                    LOGGER.log(
                            Level.WARNING,
                            "Exception sending product to "
                                    + socketSender.getHost() + ":"
                                    + socketSender.getPort(),
                            sendExceptions.get(sender));
                } else {
                    LOGGER.log(Level.WARNING, "Exception sending product "
                            + sendExceptions.get(sender));
                }
            }

            if (sendExceptions.size() < builder.getProductSenders().size()) {
                LOGGER.warning("Partial failure sending product,"
                        + " at least one sender accepted product."
                        + " Check the tracker for more information.");
                // still output built product id
                System.out.println(product.getId().toString());
                if (batchMode) {
                    // don't interrupt the batch
                    return;
                }
                // but exit with partial failure
                System.exit(CLIProductBuilder.EXIT_PARTIALLY_SENT);
            } else {
                LOGGER.severe("Total failure sending product");
                // still output built product id
                System.err.println("ERROR: " + product.getId().toString());
                if (batchMode) {
                    // don't interrupt the batch
                    return;
                }
                System.exit(CLIProductBuilder.EXIT_UNABLE_TO_SEND);
            }

        }

        // otherwise output built product id
        System.out.println(product.getId().toString());
    }

}