implemented restore & refactored code

bucket-tool

@@ -61,6 +61,22 @@ def log (message, level = 0)
   end
 end
 
+def get_verification (prompt = "Do you want to continue?")
+  while true
+    puts "#{prompt} (y/n): "
+    answer = STDIN.gets
+    answer = answer.strip()
+    case answer
+      when 'y','yes'
+        return true
+      when 'n', 'no'
+        return false
+      else
+        puts "#{answer} is not an acceptible answer"
+    end
+  end
+end
+
 # END: Helper Functions
 
 
@@ -81,7 +97,7 @@ end
 $DEBUG=false
 
 $CONFIG = Hash.new
-$CONFIG[:bucket_dir]=` sudo puppet agent --configprint clientbucketdir `.strip()
+$CONFIG[:bucket_dir]=` puppet agent --configprint clientbucketdir `.strip()
 $CONFIG[:action]=""
 $CONFIG[:search_term]=""
 $CONFIG[:log_file]=""
@@ -108,8 +124,8 @@ if not (ARGV & ['-d', '--debug']).empty?
 end
 i=0
 case ARGV[i] 
-  when 'search'
-    $CONFIG[:action]='search'
+  when 'search', 'get', 'restore'
+    $CONFIG[:action]=ARGV[i]
     log "$CONFIG[:action] was set to #{ARGV[i]}"
     log "user provided search action ARGV[i.next] == #{ARGV[i.next]}"
     if ARGV[i.next] != "" and not ARGV[i.next] =~ FLAG_REGEX
@@ -121,27 +137,9 @@ case ARGV[i]
       usage
       exit
     end
-  when 'get'
-    $CONFIG[:action] = 'get'
-    log "$CONFIG[:action] was set to #{ARGV[i]}"
-    log "user provided get action ARGV[i.next] == #{ARGV[i.next]}"
-    if ARGV[i.next] != "" and not ARGV[i.next] =~ FLAG_REGEX
-      $CONFIG[:search_term]=ARGV[i.next]
-      log "search_term was set to #{ARGV[i.next]}"
-      i+=2
-    else
-      puts "Flag[#{ARGV[i]}] : Argument[#{ARGV[i.next]}] : Either the argument was not provided or it was a flag"
-      usage
-      exit
-    end
 
-  when 'list'
-    $CONFIG[:action] = 'list'
-    log "$CONFIG[:action] was set to #{ARGV[i]}"
-    i+=1
-  
-  when 'list-files'
-    $CONFIG[:action] = 'list-files'
+  when 'list', 'list-files'
+    $CONFIG[:action] = ARGV[i]
     log "$CONFIG[:action] was set to #{ARGV[i]}"
     i+=1
 
@@ -278,6 +276,28 @@ def list_entry_files (bucket)
   puts filenames.sort.join("\n")
 end
 
+def restore_entry (bucket)
+  if bucket.entries.has_key? $CONFIG[:search_term]
+    entry = bucket.entries[$CONFIG[:search_term]]
+    if entry.filepaths.count == 1
+      filepath = entry.filepaths[0]
+      if filepath[0] != '/'
+        filepath = "/#{filepath}"
+      end
+      if get_verification "Are you sure you want to overwrite #{filepath}?" 
+        File.open(filepath,'w') do |file|
+          file.write(entry.content)
+        end
+      else
+        puts "Ok not overwriting."
+      end
+    else
+    end
+  else
+    puts "There were no entries corresponding to #{$CONFIG[:search_term]}"
+  end
+end
+
 # END: Work Functions
 
 
@@ -299,6 +319,9 @@ if __FILE__ == $0
     when 'list-files'
       list_entry_files bucket
 
+    when 'restore'
+      restore_entry bucket
+
   end
 
 end