Finished feature

bucket-tool

@@ -1,5 +1,10 @@
 #!/usr/bin/ruby
 
+# BEGIN: Requires
+
+require 'digest'
+
+# END: Requires
 
 # BEGIN: Helper Functions
 
@@ -14,15 +19,24 @@ Description:
   search the filebucket and restore from it.
 
 Actions:
-  search <term>        : Search for bucket entries matching a portion of the filepath
-  list                 : List all Bucket entries
-  list-files           : List all files/paths that have been backed up to the bucket
-  get <entry-hash>     : Get the content of a specific entry (by hash)
-  restore <entry-hash> : Restore previous state of file stored in bucket (by-hash)
+  search <term>             : Search for bucket entries matching a portion of the filepath
+  list                      : List all Bucket entries
+  list-files                : List all files/paths that have been backed up to the bucket
+  get <entry-hash>          : Get the content of a specific entry (by hash)
+  restore <value>           : Restore previous state of file stored in bucket. Value can be hash or filename/filepath
 
 Global Flags:
-  -d | --debug         : Set debug flag
-  -h | --help          : This help message
+  -d | --debug              : Set debug flag
+  -h | --help               : This help message
+
+Info Format Flags:
+  -i | --inline             : Set the info format to inline (MTIME : HASH : FILENAME)
+  -l | --long               : Set the info format to long
+                            : Entry [HASH]:
+                            :   Paths: path1,path2,...,pathn
+                            :   MTIME: YYYY-MM-DD HH:MM:SS -####
+                            :
+  -c | --csv                : Set the info format to csv ( MTIME,HASH,FILENAME1[;FILENAMEn] )
 
 Author:
   Name: Tristan Ancelet
@@ -62,6 +76,18 @@ def log (message, level = 0)
   end
 end
 
+def which(cmd)
+  #https://stackoverflow.com/questions/2108727/which-in-ruby-checking-if-program-exists-in-path-from-ruby
+  exts = ENV['PATHEXT'] ? ENV['PATHEXT'].split(';') : ['']
+  ENV['PATH'].split(File::PATH_SEPARATOR).each do |path|
+    exts.each do |ext|
+      exe = File.join(path, "#{cmd}#{ext}")
+      return exe if File.executable?(exe) && !File.directory?(exe)
+    end
+  end
+  nil
+end
+
 def get_verification (prompt = "Do you want to continue?")
   while true
     puts "#{prompt} (y/n): "
@@ -78,6 +104,54 @@ def get_verification (prompt = "Do you want to continue?")
   end
 end
 
+def get_selections (reference_array, prompt = "Which of the following do you want to select? ", options = { :multiple => false, }, &procedure)
+  ## Making clone of array since the selections were passed by reference
+  selections = reference_array.clone
+
+  def put_prompt (selections, prompt)
+    puts prompt
+    selections.each_with_index do |value,index|
+      puts "#{index} : #{value}"
+    end
+  end
+
+  if options[:multiple] == true
+    output = Array.new
+  else
+    output = ""
+  end
+
+  put_prompt selections, prompt
+
+  while true
+    choice = Integer(STDIN.gets.strip())
+    if choice.is_a? Integer
+     if  choice >= 0 and choice < selections.count 
+       if options[:multiple] == true
+         output.push(selections[choice])
+         selections.delete_at(choice)
+         if get_verification "Are you done selecting?" 
+           break
+         end
+         put_prompt selections, prompt
+       else
+         output = selections[choice]
+         break
+       end
+     else
+       puts "#{choice} is not between the values of 0 and #{selections.count}. Please try again."
+     end
+    else
+      puts "#{choice} is not a valid option. Please try again."
+    end
+  end
+  if procedure.respond_to? "call"
+    output = procedure.call(output)
+  end
+  output
+end
+
+
 # END: Helper Functions
 
 
@@ -97,11 +171,19 @@ end
 
 $DEBUG=false
 
+puppet_exe = which "puppet"
+
+if puppet_exe == nil
+  puts "The puppet utility was not found in $PATH. This utility will not be able to function"
+  exit
+end
+
 $CONFIG = Hash.new
-$CONFIG[:bucket_dir]=` puppet agent --configprint clientbucketdir `.strip()
+$CONFIG[:bucket_dir]=` #{puppet_exe} agent --configprint clientbucketdir `.strip()
 $CONFIG[:action]=""
 $CONFIG[:search_term]=""
 $CONFIG[:log_file]=""
+$CONFIG[:info_format]="inline"
 File.open('/etc/hostname') do |file|
   $HOSTNAME=file.read().strip()
 end
@@ -150,6 +232,25 @@ case ARGV[i]
     exit
 end
 
+while i < ARGV.count
+  case ARGV[i]
+    when '-c', '--csv'
+      $CONFIG[:info_format]='csv'
+      log "$CONFIG[:info_format] was set to #{$CONFIG[:info_format]}"
+      i+=1
+    when '-l', '--long'
+      $CONFIG[:info_format]='long'
+      log "$CONFIG[:info_format] was set to #{$CONFIG[:info_format]}"
+      i+=1
+    when '-i', '--inline'
+      $CONFIG[:info_format]='inline'
+      log "$CONFIG[:info_format] was set to #{$CONFIG[:info_format]}"
+      i+=1
+    else
+      i+=1
+  end
+end
+
 ## BEGIN: Checks
 
 if $CONFIG[:action] == ""
@@ -163,6 +264,7 @@ case $CONFIG[:action]
       usage
       exit
     end
+
 end
 
 ## END: Checks
@@ -177,12 +279,8 @@ class BucketEntry
   def initialize (entry_dir)
     @entry_dir = entry_dir
     @hash = File.basename(entry_dir)
-    @filepaths = Array.new
     File.open("#{entry_dir}/paths") do |file|
-      file.read().split(/\n/).each do |path|
-        log "BucketEntry[#{@hash}] adding #{path} to @filepaths"
-        @filepaths.push(path)
-      end
+      @filepaths = file.read.split(/\n/)
     end
     @mtime = File.mtime(entry_dir)
     log "BucketEntry was created from #{entry_dir}"
@@ -190,10 +288,10 @@ class BucketEntry
 
   def path_include? (path_string)
     log "BucketEntry[#{hash}] was called with #{path_string}"
-    @filepaths.each.any? {|path| path.include? path_string}
+    @filepaths.any?{|path| path.include? path_string}
   end
 
-  def infostring
+  def long_info
     "Entry [#{@hash}]:
       Paths: #{@filepaths.join(',')}
       MTIME: #{@mtime}
@@ -201,10 +299,25 @@ class BucketEntry
 "
   end
 
+  def csv_info
+    [@mtime,@hash,@filepaths.join(';')].join(',')
+  end
+
   def inline_info
     "#{@mtime} : #{@hash} : #{@filepaths.join(',')}"
   end
 
+  def info
+    case $CONFIG[:info_format]
+      when 'long'
+        long_info
+      when 'inline'
+        inline_info
+      when 'csv'
+        csv_info
+    end
+  end
+
   def content 
     log "BucketEntry[#{@hash}] getting contents"
     File.open("#{@entry_dir}/contents",'r') do |file|
@@ -223,6 +336,15 @@ class Bucket
     load_bucket
   end
 
+  def select(&proc)
+    @entries.each_value.select &proc
+  end
+
+  def any?(&proc)
+    @entries.each_value.any? &proc
+  end
+    
+
   def load_bucket
     log "Bucket[#{@bucketdir}] is loading entries"
     Dir["#{@bucketdir}/**/paths"].each.map{|path| File.dirname(path)}.each do |directory|
@@ -233,6 +355,18 @@ class Bucket
     end
     log "Bucket[#{@bucketdir}] was loaded"
   end
+
+  def filenames
+    filenames = Array.new
+    @entries.each_value do |entry|
+      entry.filepaths.each do |path|
+        if not filenames.include? path
+          filenames.push(path)
+        end
+      end
+    end
+    filenames
+  end
 end
 
 # END: Classes
@@ -242,17 +376,10 @@ end
 # BEGIN: Work Functions
 
 def search_entries_paths (bucket)
-  log "user entered"
-  bucket.entries.each_value do |entry|
-    log "checking Entry[#{entry.hash}]"
-    if entry.path_include? $CONFIG[:search_term]
-      puts entry.inline_info
-    end
-  end
+  puts bucket.select{|entry| entry.path_include? $CONFIG[:search_term]}.sort_by(&:mtime).map(&:info)
 end
 
 def get_content_of_entry_hash (bucket)
-  log "user entered"
   if bucket.entries.has_key? $CONFIG[:search_term]
     puts bucket.entries[$CONFIG[:search_term]].content
   else
@@ -262,40 +389,84 @@ def get_content_of_entry_hash (bucket)
 end
 
 def list_all_entries (bucket)
-  puts bucket.entries.each_value.each.map{|entry| entry.inline_info}.sort.join("\n")
+  puts bucket.filenames.map{|filename| bucket.select{|entry| entry.path_include? filename}.sort_by(&:mtime).map(&:info)}
 end
 
 def list_entry_files (bucket)
-  filenames = Array.new
-  bucket.entries.each_value do |entry| 
-    entry.filepaths.each do |path|
-      if not filenames.include? path
-        filenames.push(path)
+  puts bucket.filenames.sort.join("\n")
+end
+
+def get_entry_by_file (bucket, filenames)
+  entry = nil
+  if filenames.count == 1
+    filename = filenames[0]
+  else
+    filename = get_selections filenames, "Your filename matched multiple files. Please select one to restore"
+  end
+
+  entries = bucket.select{|entry| entry.path_include? filename}
+  if entries.count == 1
+    entry = entries.first
+  else
+    while true 
+      mtimes = entries.map{|entry| entry.mtime}
+      entry_mtime = get_selections(mtimes , "Which timestamp to you want to revert the file to?") 
+      entry = entries.lazy.select{|entry| entry.mtime == entry_mtime}.first
+
+      if get_verification "Do you want to see the contents of #{filename} at this time?"
+        puts entry.content
+        if get_verification "Is this the entry you want to overwrite #{filename} with?"
+          break
+        end
+      else
+        break
       end
     end
   end
-  puts filenames.sort.join("\n")
+
+  if filename[0] != '/'
+    filename = "/#{filename}"
+  end
+
+  return entry, filename
+end
+
+def get_entry_by_hash (bucket)
+  if bucket.entries.has_key? $CONFIG[:search_term]
+    entry = bucket.entries[$CONFIG[:search_term]]
+    filepath = ""
+    if entry.filepaths.count == 1
+      filepath = entry.filepaths[0]
+    else 
+      filepath = get_selections entry.filepaths, "What filepath do you wish to restore to?"
+    end
+
+    if filepath[0] != '/'
+      filepath = "/#{filepath}"
+    end
+
+    return entry, filepath
+
+  else
+    puts "There were no entries corresponding to #{$CONFIG[:search_term]}"
+  end
 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
+  entry = nil
+  if bucket.any?{ |entry| entry.path_include? $CONFIG[:search_term] }
+    filenames = bucket.filenames.select {|filename| filename.include? $CONFIG[:search_term]}
+    entry, filepath = get_entry_by_file bucket, filenames
+  else
+    entry, filepath = get_entry_by_hash bucket
+  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 "There were no entries corresponding to #{$CONFIG[:search_term]}"
+    puts "Ok not overwriting."
   end
 end