puppet-filebucket-util/bucket-tool

#!/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