Fix compilation. Add auth parameters.

adyatlov · Mar 29, 2017 · 05258d6 · aquamatthias · Mar 29, 2017 · 05258d6
1 parent 551ee1b
commit 05258d6
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 7 deletions.
diff --git a/readme.md b/readme.md
@@ -13,4 +13,7 @@
       -r,--restore-znode <arg>   the znode into which read data should be
                                  restored
 
+      -u, --user <arg>           user name, optional
+      -p, --password <arg>       password, optional
+
       -v,--verbose               enable debug output
diff --git a/src/main/java/com/d2fn/guano/DumpJob.java b/src/main/java/com/d2fn/guano/DumpJob.java
@@ -1,11 +1,14 @@
 package com.d2fn.guano;
 
-import org.apache.zookeeper.KeeperException;
 import org.apache.zookeeper.WatchedEvent;
 import org.apache.zookeeper.Watcher;
 import org.apache.zookeeper.ZooKeeper;
 import org.apache.zookeeper.data.Stat;
 
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.UnsupportedEncodingException;
 import java.util.List;
 
 /**
@@ -17,11 +20,15 @@ public class DumpJob implements Job, Watcher {
     private String zkServer;
     private String outputDir;
     private String znode;
+    private String user;
+    private String pass;
 
-    public DumpJob(String zkServer, String outputDir, String znode) {
+    public DumpJob(String zkServer, String outputDir, String znode, String user, String pass) {
         this.zkServer = zkServer;
         this.outputDir = outputDir;
         this.znode = znode;
+        this.user = user;
+        this.pass = pass;
     }
 
     public void go() {
@@ -30,13 +37,23 @@ public void go() {
         System.out.println("reading from zookeeper path: " + znode);
         System.out.println("dumping to local directory: " + outputDir);
 
+        ZooKeeper zk;
         try {
-            ZooKeeper zk = new ZooKeeper(zkServer + znode, 10000, this);
-            go(zk);
+            zk = new ZooKeeper(zkServer + znode, 10000, this);
         } catch (IOException e) {
             System.err.println("error connecting to " + zkServer);
             System.exit(1);
+            return;
         }
+        if (user != null && pass != null) {
+            try {
+                zk.addAuthInfo("digest", (user + ":" + pass).getBytes("UTF-8"));
+            } catch (UnsupportedEncodingException e) {
+                System.err.println("error processing credentials: " + e.getMessage());
+                System.exit(1);
+            }
+        }
+        go(zk);
     }
 
     private void go(ZooKeeper zk) {

diff --git a/src/main/java/com/d2fn/guano/Guano.java b/src/main/java/com/d2fn/guano/Guano.java
@@ -33,18 +33,25 @@ public void buildJob() {
         }
 
         String server = cmd.getOptionValue("s");
+        String user = null;
+        String pass = null;
+
+        if (cmd.hasOption("u") && cmd.hasOption("p")) {
+            user = cmd.getOptionValue("u");
+            pass = cmd.getOptionValue("p");
+        }
 
         // dump?
         if(cmd.hasOption("d") && cmd.hasOption("o")) {
             String znode = cmd.getOptionValue("d");
             String outputDir = cmd.getOptionValue("o");
-            job = new DumpJob(server, outputDir, znode);
+            job = new DumpJob(server, outputDir, znode, user, pass);
         }
         // restore
         else if(cmd.hasOption("r") && cmd.hasOption("i")) {
             String znode = cmd.getOptionValue("r");
             String inputDir = cmd.getOptionValue("i");
-            job = new RestoreJob(server, znode, inputDir);
+            job = new RestoreJob(server, znode, inputDir, user, pass);
         }
         else {
             usage(options);
@@ -84,6 +91,8 @@ public static Options initOptions() {
         options.addOption("o", "output-dir",        true,  "the output directory to which znode information should be written (must be a normal, empty directory)");
         options.addOption("i", "input-dir",         true,  "the input directory from which znode information should be read");
         options.addOption("v", "verbose",           false, "enable debug output");
+        options.addOption("u", "user",           true, "user name, optional");
+        options.addOption("p", "password",           true, "password, optional");
         return options;
     }
 }
diff --git a/src/main/java/com/d2fn/guano/RestoreJob.java b/src/main/java/com/d2fn/guano/RestoreJob.java
@@ -16,13 +16,17 @@ public class RestoreJob implements Job, Watcher, FSVisitor {
     private String zkServer;
     private String znode;
     private String inputDir;
+    private String user;
+    private String pass;
 
     private ZooKeeper zk;
 
-    public RestoreJob(String zkServer, String znode, String inputDir) {
+    public RestoreJob(String zkServer, String znode, String inputDir, String user, String pass) {
         this.zkServer = zkServer;
         this.znode = znode;
         this.inputDir = inputDir;
+        this.user = user;
+        this.pass = pass;
     }
 
     public void go() {
@@ -33,6 +37,9 @@ public void go() {
 
         try {
             zk = new ZooKeeper(zkServer + znode, 10000, this);
+            if (user != null && pass != null) {
+                zk.addAuthInfo("digest", (user + ":" + pass).getBytes("UTF-8"));
+            }
             while(!zk.getState().isConnected()) {
                 System.out.println("connecting to " + zkServer + " with chroot " + znode);
                 Thread.sleep(1000L);