apt automation tool for use in puppet

apt-tool.rb

#!/usr/bin/env ruby

require 'open3'
require 'open-uri'
require 'getoptlong'
require 'digest/md5'

class Apt

  def initialize
    @source = nil
    @update = false
  end

  #------------------------------------------------------------------------------#
  ## PPA operations
  #------------------------------------------------------------------------------#
  
  def installed_ppas
    retv = Array.new
    Dir.glob('/etc/apt/sources.list.d/*.list') {|fname|
      File.open(fname, 'r').each_line {|line|
        if line =~ /^.*(ht|f)[tps]+:\/\/ppa.launchpad.net.*/
          ppa = "ppa:#{line.split()[1].split('net/')[1].split('/ubun')[0]}"
          if not retv.include? ppa
            retv.push(ppa)
          end
        end
      }
    }
    retv
  end
  
  def installed_ppa?(ppa=nil)
    if ppa
      if installed_ppas.include? ppa
        return true
      else
        return false
      end
    end
  end
  
  def install_ppa(ppa=nil)
    if ppa
      if ppa =~ /^ppa:.*\/.*$/ and not installed_ppa? ppa
        if run_cmd("add-apt-repository -y #{ppa}") == true
          return update(true)
        end
      end
    end
    return false
  end

  def remove_ppa(ppa=nil)
    if ppa
      if ppa =~ /^ppa:.*\/.*$/ and installed_ppa? ppa
        if run_cmd("add-apt-repository -y -r #{ppa}") == true
          return update(true)
        end
      end
      return false
    end
  end

  #------------------------------------------------------------------------------#
  ## Non-PPA repo sources
  #------------------------------------------------------------------------------#

  def installed_srcs
    retv = Array.new
    Dir.glob('/etc/apt/sources.list.d/*.list') {|fname|
      File.open(fname, 'r').each_line {|line|
        if line =~ /^deb\ .*/
          tmp = line.split
          src_uri = "#{tmp[1]}/dists/#{tmp[2]}/#{tmp[3]}/"
          if not retv.include? [src_uri, fname]
            retv.push([src_uri, fname])
          end
        end
      }
    }
    return retv
  end
  
  def installed_src?(src=nil)
    if src
      installed_srcs.each {|itm|
        if itm.include? src
          return itm
        end
      }
    end
    return false
  end
  
  def install_src(src=nil)
    if src and src =~ /^(ht|f)[tps]+:\/\/.*\/dists\/.*$/
      if installed_src?(src) == false
        fname   = "/etc/apt/sources.list.d/#{src.split('//')[1].split('/')[0]}.list"
        while File.exist? fname
          fname = "/etc/apt/sources.list.d/#{src.split('//')[1].split('/')[0]}-#{rand(100..999)}.list"
        end
        tmp     = src.split('/dists/')
        File.open(fname, 'w+') {|fp| fp.write "deb #{tmp[0]} #{tmp[1].split('/')[0]} #{tmp[1].split('/')[1]}\n" }
        return update(true)
      end
    end
    return false
  end

  def remove_src(src=nil)
    if src and src =~ /^(ht|f)[tps]+:\/\/.*\/dists\/.*$/
      itm = installed_src? src
      if itm != false and itm.class == Array
        File.delete itm[1]
        return update(true)
      end
    end
    return false
  end

  #------------------------------------------------------------------------------#
  ## Manual key(from uri) management
  #------------------------------------------------------------------------------#

  def installed_keys
    retv     = Array.new
    key_id   = nil
    key_str  = nil
    `/usr/bin/apt-key list`.split("\n").each {|l|
      key_id  = l.split[1].split('/')[1]       if l =~ /^pub\ .*$/
      key_str = l.split.slice(1..-1).join(' ') if l =~ /^uid\ .*$/
      if not [key_id, key_str].include? nil
        retv.push("#{key_id}:#{key_str}")
        key_id  = nil
        key_str = nil
      end
    }
    return retv
  end

  def installed_key?(key=nil)
    if key
      installed_keys.each {|ln|
        tmp = ln.split(':')
        if key == tmp[0] or key == tmp[1]
          return true
        end
      }
      return false
    end
  end

  def install_key(key=nil)
    if key 
      if key =~ /^(ht|f)[tps]+:\/\/.*$/ and installed_key? key
        return run_cmd("wget -q -O- #{key} | apt-key add -")
      elsif key =~ /^[\.\/a-zA-Z]+.*$/ and not installed_key? key
        return run_cmd("apt-key add #{key}")
      else
        puts 'You have supplied an invalid key specification.'
        exit 255
      end
    else
      return false
    end
  end

  def remove_key(key=nil)
    if key and key =~ /[A-Z0-9]{,8}/
      return run_cmd("apt-key del #{key}")
    else
      return false
    end
  end

  #------------------------------------------------------------------------------#
  ## apt-get update
  #------------------------------------------------------------------------------#

  def update(update=false)
    if update
      return run_cmd('/usr/bin/apt-get update') if update
    end
  end

  #------------------------------------------------------------------------------#
  ## private utility methods
  #------------------------------------------------------------------------------#

  private
  def run_cmd(cmd=nil)
    lines = Array.new
    if cmd
      `#{cmd}`.split("\n").each {|ln|
        lines.push(ln)
      }
      if $?.exitstatus != 0
        if lines.size > 10
          lines[-10..-1].each {|m|
            puts m
          }
        else
          puts lines
        end
        return false
      else
        return true
      end
    end
  end

end

def show_help
  puts <<-EOF

#{File.basename(__FILE__)} [operation] [type] <argument>

--help,   -h            Show this help screen.
--add,    -A            Set the action to 'add'
--check,  -C            Set the action to 'check'
--remove, -R            Set the action to 'remove'
--source, -s <URI>      Set the source to <URI> (ex: http://mirrors.us.kernel.org/ubuntu/dists/precise/main/)
--ppa,    -p <PPA>      Set the source to <PPA> (ex: ppa:devteam-p/production)
--key,    -k <KEY>      Set the source to <KEY> (keyID, or keyDESCR will work)

EOF
end

#------------------------------------------------------------------------------#
## Option handling
#------------------------------------------------------------------------------#

if ARGV.size < 1
  show_help
  exit
end

opts = GetoptLong.new([ '--help',   '-h', GetoptLong::NO_ARGUMENT       ],
                      [ '--add',    '-A', GetoptLong::NO_ARGUMENT       ],
                      [ '--remove', '-R', GetoptLong::NO_ARGUMENT       ],
                      [ '--check',  '-C', GetoptLong::NO_ARGUMENT       ],
                      [ '--source', '-s', GetoptLong::REQUIRED_ARGUMENT ],
                      [ '--ppa',    '-p', GetoptLong::REQUIRED_ARGUMENT ],
                      [ '--key',    '-k', GetoptLong::REQUIRED_ARGUMENT ])

$apt   = Apt.new
action = nil
type   = nil
source = nil
begin
  opts.each {|opt, arg|
    case opt
      when '--help'
      show_help
      exit 0
      when '--add'
      action = :add
      when '--remove'
      action = :remove
      when '--check'
      action = :check
      when '--source'
      type   = :source
      source = arg
      when '--ppa'
      type   = :ppa
      source = arg
      when '--key'
      type   = :key
      source = arg
    end
  }
rescue GetoptLong::MissingArgument => m
  exit 255
rescue GetoptLong::InvalidOption => m
  exit 255
rescue ArgumentError => m
  puts "Error: #{m}"
  exit 255
end

#------------------------------------------------------------------------------#
## And main logic
#------------------------------------------------------------------------------#

case action
  when :check
  case type
    when :source
    exit 0 if $apt.installed_src? source
    exit 1 if not $apt.installed_src? source
    when :ppa
    exit 0 if $apt.installed_ppa? source
    exit 1 if not $apt.installed_ppa? source
    when :key
    exit 0 if $apt.installed_key? source
    exit 1 if not $apt.installed_key? source
  end
  when :add
  case type
    when :source
    exit 0 if $apt.install_src source
    exit 1 if not $apt.install_src source
    when :ppa
    exit 0 if $apt.install_ppa source
    exit 1 if not $apt.install_ppa source
    when :key
    exit 0 if $apt.install_key source
    exit 1 if not $apt.install_key source
  end
  when :remove
  case type
    when :source
    exit 0 if $apt.remove_src source
    exit 1 if not $apt.remove_src source
    when :ppa
    exit 0 if $apt.remove_ppa source
    exit 1 if not $apt.remove_ppa source
    when :key
    exit 0 if $apt.remove_key source
    exit 1 if not $apt.remove_key source
  end
end