--- /dev/null
+require File.expand_path('../external_iterator', __FILE__)
+require File.expand_path('../ini_file/section', __FILE__)
+
+module Puppet
+module Util
+ class IniFile
+
+ def initialize(path, key_val_separator = ' = ', section_prefix = '[', section_suffix = ']')
+
+ k_v_s = key_val_separator.strip
+
+ @section_prefix = section_prefix
+ @section_suffix = section_suffix
+
+ @@SECTION_REGEX = section_regex
+ @@SETTING_REGEX = /^(\s*)([^#;\s]|[^#;\s].*?[^\s#{k_v_s}])(\s*#{k_v_s}\s*)(.*)\s*$/
+ @@COMMENTED_SETTING_REGEX = /^(\s*)[#;]+(\s*)(.*?[^\s#{k_v_s}])(\s*#{k_v_s}[ \t]*)(.*)\s*$/
+
+ @path = path
+ @key_val_separator = key_val_separator
+ @section_names = []
+ @sections_hash = {}
+ if File.file?(@path)
+ parse_file
+ end
+ end
+
+ def section_regex
+ # Only put in prefix/suffix if they exist
+ # Also, if the prefix is '', the negated
+ # set match should be a match all instead.
+ r_string = '^\s*'
+ r_string += Regexp.escape(@section_prefix)
+ r_string += '('
+ if @section_prefix != ''
+ r_string += '[^'
+ r_string += Regexp.escape(@section_prefix)
+ r_string += ']'
+ else
+ r_string += '.'
+ end
+ r_string += '*)'
+ r_string += Regexp.escape(@section_suffix)
+ r_string += '\s*$'
+ /#{r_string}/
+ end
+
+ def section_names
+ @section_names
+ end
+
+ def get_settings(section_name)
+ section = @sections_hash[section_name]
+ section.setting_names.inject({}) do |result, setting|
+ result[setting] = section.get_value(setting)
+ result
+ end
+ end
+
+ def get_value(section_name, setting)
+ if (@sections_hash.has_key?(section_name))
+ @sections_hash[section_name].get_value(setting)
+ end
+ end
+
+ def set_value(section_name, setting, value)
+ unless (@sections_hash.has_key?(section_name))
+ add_section(Section.new(section_name, nil, nil, nil, nil))
+ end
+
+ section = @sections_hash[section_name]
+
+ if (section.has_existing_setting?(setting))
+ update_line(section, setting, value)
+ section.update_existing_setting(setting, value)
+ elsif result = find_commented_setting(section, setting)
+ # So, this stanza is a bit of a hack. What we're trying
+ # to do here is this: for settings that don't already
+ # exist, we want to take a quick peek to see if there
+ # is a commented-out version of them in the section.
+ # If so, we'd prefer to add the setting directly after
+ # the commented line, rather than at the end of the section.
+
+ # If we get here then we found a commented line, so we
+ # call "insert_inline_setting_line" to update the lines array
+ insert_inline_setting_line(result, section, setting, value)
+
+ # Then, we need to tell the setting object that we hacked
+ # in an inline setting
+ section.insert_inline_setting(setting, value)
+
+ # Finally, we need to update all of the start/end line
+ # numbers for all of the sections *after* the one that
+ # was modified.
+ section_index = @section_names.index(section_name)
+ increment_section_line_numbers(section_index + 1)
+ else
+ section.set_additional_setting(setting, value)
+ end
+ end
+
+ def remove_setting(section_name, setting)
+ section = @sections_hash[section_name]
+ if (section.has_existing_setting?(setting))
+ # If the setting is found, we have some work to do.
+ # First, we remove the line from our array of lines:
+ remove_line(section, setting)
+
+ # Then, we need to tell the setting object to remove
+ # the setting from its state:
+ section.remove_existing_setting(setting)
+
+ # Finally, we need to update all of the start/end line
+ # numbers for all of the sections *after* the one that
+ # was modified.
+ section_index = @section_names.index(section_name)
+ decrement_section_line_numbers(section_index + 1)
+ end
+ end
+
+ def save
+ File.open(@path, 'w') do |fh|
+
+ @section_names.each_index do |index|
+ name = @section_names[index]
+
+ section = @sections_hash[name]
+
+ # We need a buffer to cache lines that are only whitespace
+ whitespace_buffer = []
+
+ if (section.is_new_section?) && (! section.is_global?)
+ fh.puts("\n#{@section_prefix}#{section.name}#{@section_suffix}")
+ end
+
+ if ! section.is_new_section?
+ # write all of the pre-existing settings
+ (section.start_line..section.end_line).each do |line_num|
+ line = lines[line_num]
+
+ # We buffer any lines that are only whitespace so that
+ # if they are at the end of a section, we can insert
+ # any new settings *before* the final chunk of whitespace
+ # lines.
+ if (line =~ /^\s*$/)
+ whitespace_buffer << line
+ else
+ # If we get here, we've found a non-whitespace line.
+ # We'll flush any cached whitespace lines before we
+ # write it.
+ flush_buffer_to_file(whitespace_buffer, fh)
+ fh.puts(line)
+ end
+ end
+ end
+
+ # write new settings, if there are any
+ section.additional_settings.each_pair do |key, value|
+ fh.puts("#{' ' * (section.indentation || 0)}#{key}#{@key_val_separator}#{value}")
+ end
+
+ if (whitespace_buffer.length > 0)
+ flush_buffer_to_file(whitespace_buffer, fh)
+ else
+ # We get here if there were no blank lines at the end of the
+ # section.
+ #
+ # If we are adding a new section with a new setting,
+ # and if there are more sections that come after this one,
+ # we'll write one blank line just so that there is a little
+ # whitespace between the sections.
+ #if (section.end_line.nil? &&
+ if (section.is_new_section? &&
+ (section.additional_settings.length > 0) &&
+ (index < @section_names.length - 1))
+ fh.puts("")
+ end
+ end
+
+ end
+ end
+ end
+
+
+ private
+ def add_section(section)
+ @sections_hash[section.name] = section
+ @section_names << section.name
+ end
+
+ def parse_file
+ line_iter = create_line_iter
+
+ # We always create a "global" section at the beginning of the file, for
+ # anything that appears before the first named section.
+ section = read_section('', 0, line_iter)
+ add_section(section)
+ line, line_num = line_iter.next
+
+ while line
+ if (match = @@SECTION_REGEX.match(line))
+ section = read_section(match[1], line_num, line_iter)
+ add_section(section)
+ end
+ line, line_num = line_iter.next
+ end
+ end
+
+ def read_section(name, start_line, line_iter)
+ settings = {}
+ end_line_num = nil
+ min_indentation = nil
+ while true
+ line, line_num = line_iter.peek
+ if (line_num.nil? or match = @@SECTION_REGEX.match(line))
+ return Section.new(name, start_line, end_line_num, settings, min_indentation)
+ elsif (match = @@SETTING_REGEX.match(line))
+ settings[match[2]] = match[4]
+ indentation = match[1].length
+ min_indentation = [indentation, min_indentation || indentation].min
+ end
+ end_line_num = line_num
+ line_iter.next
+ end
+ end
+
+ def update_line(section, setting, value)
+ (section.start_line..section.end_line).each do |line_num|
+ if (match = @@SETTING_REGEX.match(lines[line_num]))
+ if (match[2] == setting)
+ lines[line_num] = "#{match[1]}#{match[2]}#{match[3]}#{value}"
+ end
+ end
+ end
+ end
+
+ def remove_line(section, setting)
+ (section.start_line..section.end_line).each do |line_num|
+ if (match = @@SETTING_REGEX.match(lines[line_num]))
+ if (match[2] == setting)
+ lines.delete_at(line_num)
+ end
+ end
+ end
+ end
+
+ def create_line_iter
+ ExternalIterator.new(lines)
+ end
+
+ def lines
+ @lines ||= IniFile.readlines(@path)
+ end
+
+ # This is mostly here because it makes testing easier--we don't have
+ # to try to stub any methods on File.
+ def self.readlines(path)
+ # If this type is ever used with very large files, we should
+ # write this in a different way, using a temp
+ # file; for now assuming that this type is only used on
+ # small-ish config files that can fit into memory without
+ # too much trouble.
+ File.readlines(path)
+ end
+
+ # This utility method scans through the lines for a section looking for
+ # commented-out versions of a setting. It returns `nil` if it doesn't
+ # find one. If it does find one, then it returns a hash containing
+ # two keys:
+ #
+ # :line_num - the line number that contains the commented version
+ # of the setting
+ # :match - the ruby regular expression match object, which can
+ # be used to mimic the whitespace from the comment line
+ def find_commented_setting(section, setting)
+ return nil if section.is_new_section?
+ (section.start_line..section.end_line).each do |line_num|
+ if (match = @@COMMENTED_SETTING_REGEX.match(lines[line_num]))
+ if (match[3] == setting)
+ return { :match => match, :line_num => line_num }
+ end
+ end
+ end
+ nil
+ end
+
+ # This utility method is for inserting a line into the existing
+ # lines array. The `result` argument is expected to be in the
+ # format of the return value of `find_commented_setting`.
+ def insert_inline_setting_line(result, section, setting, value)
+ line_num = result[:line_num]
+ match = result[:match]
+ lines.insert(line_num + 1, "#{' ' * (section.indentation || 0 )}#{setting}#{match[4]}#{value}")
+ end
+
+ # Utility method; given a section index (index into the @section_names
+ # array), decrement the start/end line numbers for that section and all
+ # all of the other sections that appear *after* the specified section.
+ def decrement_section_line_numbers(section_index)
+ @section_names[section_index..(@section_names.length - 1)].each do |name|
+ section = @sections_hash[name]
+ section.decrement_line_nums
+ end
+ end
+
+ # Utility method; given a section index (index into the @section_names
+ # array), increment the start/end line numbers for that section and all
+ # all of the other sections that appear *after* the specified section.
+ def increment_section_line_numbers(section_index)
+ @section_names[section_index..(@section_names.length - 1)].each do |name|
+ section = @sections_hash[name]
+ section.increment_line_nums
+ end
+ end
+
+
+ def flush_buffer_to_file(buffer, fh)
+ if buffer.length > 0
+ buffer.each { |l| fh.puts(l) }
+ buffer.clear
+ end
+ end
+
+ end
+end
+end