initial commit

bucket-tool

@@ -0,0 +1,294 @@
+#!/usr/bin/ruby
+
+# BEGIN: Pre-Run Check
+
+if ENV["USER"] != 'root'
+  puts "This script should only be run by root (permissions issues). Please rerun it as root or prepend \"sudo\""
+  exit
+end
+
+# END: Pre-Run Check
+
+
+
+# BEGIN: Variables
+
+$DEBUG=false
+
+$CONFIG = Hash.new
+$CONFIG[:bucket_dir]=` sudo puppet agent --configprint clientbucketdir `.strip()
+$CONFIG[:action]=""
+$CONFIG[:search_term]=""
+$CONFIG[:log_file]=""
+File.open('/etc/hostname') do |file|
+  $HOSTNAME=file.read().strip()
+end
+FLAG_REGEX=/\-+\S+/
+# END: Variables
+
+
+
+# BEGIN: Helper Functions
+
+def usage
+  puts "#{__FILE__} [ACTION <arg>] [<flags>]
+
+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-md5>     : Get the content of a specific entry (by md5)
+  restore <entry-md5> : Restore entry to 
+
+Global Flags:
+  -d | --debug        : Set debug flag
+"
+end
+
+def log (message, level = 0)
+  if $DEBUG == true or $CONFIG[:log_file] != ""
+    if message == ""
+      puts "log was called without providing a message"
+      exit
+    end
+
+    case level
+      when 0
+        level="INFO"
+      when 1
+        level="WARN"
+      when 2
+        level="CRIT"
+      else
+        level="UNDF"
+    end
+    datestamp=Time.now
+    log_message="#{datestamp} : #{$HOSTNAME} : #{level} : #{caller[0]} : #{message}"
+
+    if $CONFIG[:log_file] != ""
+      File.open($CONFIG[:log_file],'a') do |file|
+        file.write("#{log_message}\n")
+      end
+    else
+      puts log_message
+    end
+  end
+end
+
+# END: Helper Functions
+
+
+
+# BEGIN: Handle CLI Args
+
+if ARGV.count == 0
+  puts "No arguments were provided"
+  usage
+  exit
+end
+
+if not (ARGV & ['-d', '--debug']).empty?
+  $DEBUG=true
+end
+
+i=0
+case ARGV[i] 
+  when 'search'
+    $CONFIG[:action]='search'
+    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
+      $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 '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'
+    log "$CONFIG[:action] was set to #{ARGV[i]}"
+    i+=1
+
+  else
+    puts "#{ARGV[i]} is not a valid action. Please make sure you use a valid action as the first argument of the script"
+    usage
+    exit
+end
+
+## BEGIN: Checks
+
+if $CONFIG[:action] == ""
+  puts "Action was not provided"
+end
+
+case $CONFIG[:action]
+  when 'search', 'get'
+    if $CONFIG[:search_term] == ""
+      puts "Search Term was not provided"
+      usage
+      exit
+    end
+end
+
+## END: Checks
+
+# END: Handle CLI Args
+
+
+# BEGIN: Classes
+
+class BucketEntry
+  attr_reader :md5, :filepaths, :mtime
+  def initialize (entry_dir)
+    @entry_dir = entry_dir
+    @md5 = File.basename(entry_dir)
+    @filepaths = Array.new
+    File.open("#{entry_dir}/paths") do |file|
+      file.read().split(/\n/).each do |path|
+        log "BucketEntry[#{@md5}] adding #{path} to @filepaths"
+        @filepaths.push(path)
+      end
+    end
+    @mtime = File.mtime(entry_dir)
+    log "BucketEntry was created from #{entry_dir}"
+  end
+
+  def path_include? (path_string)
+    log "BucketEntry[#{md5}] was called with #{path_string}"
+    @filepaths.each.any? {|path| path.include? path_string}
+  end
+
+  def infostring
+    "Entry [#{@md5}]:
+      Paths: #{@filepaths.join(',')}
+      MTIME: #{@mtime}
+
+"
+  end
+
+  def inline_info
+    "#{@mtime} : #{@md5} : #{@filepaths.join(',')}"
+  end
+
+  def content 
+    log "BucketEntry[#{@md5}] getting contents"
+    File.open("#{@entry_dir}/contents",'r') do |file|
+      file.read()
+    end
+  end
+
+end
+
+class Bucket
+  attr_reader :bucketdir, :entries
+  def initialize (clientbucketdir)
+    log "Bucket is being created from #{clientbucketdir}"
+    @bucketdir = clientbucketdir
+    @entries = Array.new
+    load_bucket
+  end
+
+  def load_bucket
+    log "Bucket[#{@bucketdir}] is loading entries"
+    Dir["#{@bucketdir}/**/paths"].each.map{|path| File.dirname(path)}.each do |directory|
+      log "\"#{directory}\" was grabbed from bucket directory. Making new BucketEntry"
+      entry = BucketEntry.new(directory)
+      @entries.push(entry)
+      log "BucketEntry[#{entry.md5}] was added to @entries Size=#{@entries.count()}"
+    end
+    log "Bucket[#{@bucketdir}] was loaded"
+  end
+end
+
+# END: Classes
+
+
+
+# BEGIN: Work Functions
+
+def search_entries_paths (bucket)
+  log "user entered"
+  bucket.entries.each do |entry|
+    log "checking Entry[#{entry.md5}]"
+    if entry.path_include? $CONFIG[:search_term]
+      puts entry.inline_info
+    end
+  end
+end
+
+def get_content_of_entry_md5 (bucket)
+  log "user entered"
+  bucket.entries.each do |entry|
+    log "checking Entry[#{entry.md5}]"
+    if entry.md5 == $CONFIG[:search_term]
+      log "BucketEntry[#{entry.md5}] Matched. Getting contents"
+      puts entry.content
+      exit
+    end
+  end
+end
+
+def list_all_entries (bucket)
+  puts bucket.entries.each.map{|entry| entry.inline_info}.sort.join("\n")
+end
+
+def list_entry_files (bucket)
+  filenames = Array.new
+  bucket.entries.each do |entry| 
+    entry.filepaths.each do |path|
+      if not filenames.include? path
+        filenames.push(path)
+      end
+    end
+  end
+  puts filenames.sort.join("\n")
+end
+
+# END: Work Functions
+
+
+# BEGIN: Work
+
+if __FILE__ == $0
+  bucket = Bucket.new($CONFIG[:bucket_dir])
+
+  case $CONFIG[:action]
+    when 'search'
+      search_entries_paths bucket
+
+    when 'get'
+      get_content_of_entry_md5 bucket
+
+    when 'list'
+      list_all_entries bucket
+
+    when 'list-files'
+      list_entry_files bucket
+
+  end
+
+end
+# END: Work
+
+