require 'chef/resource'
require 'chef/provider'
require 'chef/mixin/command'
require 'chef/log'
 
class Chef
  class Resource
    class Ban < Chef::Resource      
      if Chef::VERSION == '0.6.0.2'
        def initialize(name, collection = nil, node = nil)
          super(name, collection, node)
          init(name)
        end
      else
        def initialize(name, run_context = nil)
          super(name, run_context)
          init(name)
        end
      end
      
      def init(name)        
        @resource_name = :ban
        @action = :update
        @allowed_actions.push(:update)
        
        @ip = {}
      end

      def ip(address = nil, options = {})
        valid_option_keys = [:ports]
        invalid_keys = options.keys - valid_option_keys
        
        if address
          unless address && address.is_a?(String)
            raise ArgumentError, 'IP Address must be a string, eg: "219.13.1.98"'
          end
        
          unless invalid_keys.empty?
            raise ArgumentError, "The following are invalid options: #{invalid_keys.join(', ')}"
          end
        
          @ip.merge!(address => options)
        end
        
        set_or_return(:ip, @ip, :kind_of => [Hash])
      end
    end
  end
  
  class Provider
    class Ban < Chef::Provider
      include Chef::Mixin::Command
                  
      def load_current_resource
        true
      end
      
      def add_rule(ip, port = nil)
        # default flags
        flags = { 
          '-A' => 'INPUT', 
          '-s' => ip, 
          '-j' => 'DROP', 
          '-m' => 'comment', 
          '--comment' => 'chef' 
        }
        
        # port flags
        flags.merge!('-p' => 'tcp', '--dport' => port) if port
        
        # execute command
        iptables(flags)
      end
      
      def iptables(flags = {})
        command = flags.inject('iptables'){|m,(k,v)| m << " #{k} #{v}"}
        Chef::Log.info "Running: #{command}"
        run_command(:command => command)
      end
      
      def action_update
        # remove rules previously generated by ban recipe
        command = "for i in $(iptables --list INPUT --line-numbers | awk '/chef/{print $1}' | sort -rn); do iptables -D INPUT $i; done"
        run_command(:command => command)
        
        # write rules
        @new_resource.ip.each_pair do |ip, options|
          ports = (options[:ports] || []).select{|p| p.is_a?(Integer)}
          ports.empty? ? add_rule(ip) : ports.each {|port| add_rule(ip, port)}
        end
        
        # save rules
        run_command(:command => "/etc/init.d/iptables save")
      end
    end
  end
end

Chef::Platform.platforms[:default].merge!(:ban => Chef::Provider::Ban)