Puppet Manifest Cleanup

gistfile1.rb

#!/usr/bin/ruby
require 'parser-emit'

class Puppet::Parser::AST::Definition
  attr_accessor :name
end

class Manifest
  class << self
    #include Enumerable

    def each(&block)
      return @manifests.each(&block)
    end

    def <<(manifest)
      @manifests ||= []
      @manifests << manifest
    end
  end

  attr_accessor :name
  alias_method(:to_s, :name)

  def initialize(module_name, args)
    @module = module_name
    @rawfile = ''
    if args.has_key? :file
      @filename = args[:file]
      @name = should_contain
      load_ast_from_file
    elsif args.has_key? :name
      @name = args[:name]

      if @name == @module
        @filename = "manifests/init.pp"
      else
        @filename = "manifests/#{@name.split(/::/)[1..-1].join('/')}.pp"
      end

      raise Exception.new("No rawfile provided") unless args.has_key? :rawfile

      @rawfile = args[:rawfile]
    end

    Manifest << self
  end

  def load_ast_from_file
    @rawfile = File.open(@filename){|f| f.read}
    @ast = ast_for_str(@rawfile)
  end

  def ast_for_str(str)
    parser = Puppet::Parser::Parser.new(Puppet::Node::Environment.new)
    parser.parse(str)
  end

  def should_contain
    if @filename == 'manifests/init.pp'
      @module
    else
      parts = @filename.split(/\//)
      parts = parts[1, parts.size]
      parts.last.sub! /\.pp\Z/, ''
      "#{@module}::#{parts.join('::')}"
    end
  end

  def move_constructs
    # finds all classes and defines that do not belong in this
    # manifest.  for each found, create a new manifest and put it
    # there, and detach the original from the ast
    process_subtree(@ast)
  end

  def process_subtree(tree)
    if tree.instance_of? Puppet::Parser::AST::ASTArray
      tree.children.delete_if do |t|
        process_subtree(t) if t.instance_of? Puppet::Parser::AST::Hostclass
        if (t.instance_of?(Puppet::Parser::AST::Hostclass) ||
            t.instance_of?(Puppet::Parser::AST::Definition))
          if should_contain != t.name
            do_move(t)
          end
        end
      end
    elsif tree.instance_of? Puppet::Parser::AST::Hostclass
      process_subtree(tree.code) if tree.code
    end
  end

  def block_for_position(position, string)
    cursor = position
    stacks = { "{" => [], '(' => [], '"' => [], "'" => [], '\\' => [], '#' => [] }
    escape = false
    in_body = false

    string[cursor...-1].each_char do |c|
      begin
        case c
        when '{'
          next unless stacks['"'].empty? && stacks["'"].empty? && stacks['#'].empty?
          stacks[c].push c
          if in_body == false && stacks['('].empty?
            in_body = true
          end
        when '}'
          next unless stacks['"'].empty? && stacks["'"].empty? && stacks['#'].empty?
          stacks['{'].pop
          if stacks['{'].empty? && in_body
            return string[position..cursor]
          end
        when '('
          next unless stacks['"'].empty? && stacks["'"].empty? && stacks['#'].empty?
          stacks[c].push c
        when ')'
          next unless stacks['"'].empty? && stacks["'"].empty? && stacks['#'].empty?
          stacks['('].pop
        when '"'
          next unless stacks["'"].empty? && stacks['#'].empty?
          if stacks[c].empty?
            stacks[c].push c
          elsif stacks['\\'].empty?
            stacks[c].pop
          else
            stacks['\\'] = []
          end
        when "'"
          next unless stacks['"'].empty? && stacks['#'].empty?
          if stacks[c].empty?
            stacks[c].push c
          elsif stacks['\\'].empty?
            stacks[c].pop
          else
            stacks['\\'] = []
          end
        when '\\'
          if stacks[c].empty?
            stacks[c].push c
          else
            stacks[c] = []
          end
        when '#'
          next unless stacks['"'].empty? && stacks["'"].empty?
          stacks[c].push c
        when "\n"
          stacks['#'] = []
        else
          stacks['\\'] = []
        end
      ensure
        cursor += 1
      end
    end

    raise Exception.new "should not fall through"
  end

  def find_construct_in_string(type, name, string)
    return false if name.empty? or string.empty?

    # simple case where name is already correct
    position = string.index(Regexp.new("#{type}\\s+#{name}[\\s({]"))
    return position if position

    # well, then it must be nested, progressively increase scope
    parts = name.split('::')
    return false if parts.size == 1 # not nested... just not here
    removed = [parts.pop]
    base_position = nil

    while (parts.size >= 1)
      if found = find_construct_in_string("class", parts.join('::'), string)
        base_position = found
        break
      else
        removed.insert(0, parts.pop)
      end
    end

    nested_position = find_construct_in_string(type, removed.join('::'),
                                               block_for_position(base_position, string))

    if nested_position
      base_position + nested_position
    else
      false
    end
  end

  def do_move(tree)
    if tree.instance_of? Puppet::Parser::AST::Hostclass
      type = 'class'
    elsif tree.instance_of? Puppet::Parser::AST::Definition
      type = 'define'
    else
      raise Exception.new "Unknown tree type"
    end

    start_pos = find_construct_in_string(type, tree.name, @rawfile)
    new_manifest = block_for_position(start_pos, @rawfile)

    @rawfile = @rawfile[0...start_pos] + \
               "# The #{type} '#{tree.name}' was relocated for autoloading\n" + \
               @rawfile[(start_pos+new_manifest.size+1)..-1]

    # qualify unqualified things
    if tree.name.match(Regexp.new("^#{@module}::"))
      new_name = tree.name
    else
      new_name = "#{@module}::#{tree.name}"
    end

    munge_name!(new_manifest, type, new_name)
    Manifest.new(@module, :name => new_name, :rawfile => new_manifest)
  end

  def munge_name!(string, type, name)
    string.sub!(Regexp.new("#{type} " + '[-\w:]+'), "#{type} #{name}")
  end

  def sync!
    `mkdir -p #{File.dirname(@filename)}` unless File.directory?(File.dirname(@filename))
    File.open(@filename, 'w'){|f| f.puts @rawfile}
  end
end

def remove_imports
  `find -wholename '*/manifests/*.pp' | xargs sed -ri '/^\s*import\s+/d'`
end

unless File.directory?('./manifests')
  STDERR.puts "#{File.expand_path('.')} does not appear to be the root of a module"
  exit
end

module_name = File.basename(File.expand_path('.'))
manifest_files = `find manifests -name '*.pp'`.split /\n/

remove_imports

manifest_files.each {|m| Manifest.new module_name, :file => m}

Manifest.each do |manifest|
  manifest.move_constructs
  manifest.sync!
end