OpenNebula

class OpenNebulaFirewall < OpenNebulaNetwork

    DRIVER = "fw"

    XPATH_FILTER =  "TEMPLATE/NIC[FW_IN|FW_OUT]"

    def initialize(vm, deploy_id = nil, hypervisor = nil)

        super(vm,XPATH_FILTER,deploy_id,hypervisor)

        @locking = true

    end

    def activate

        lock

        vm_id =  @vm['ID']

        process do |nic|

            tap = nic[:tap]

            ip  = nic[:ip]

            if tap

                if rules = nic[:fw_in]

                    chain = "one-#{vm_id}-#{ip}-in"

                    nic_rules = Array.new

                    nic_rules << new_chain(chain)

                    nic_rules << tap_to_chain(tap, chain, "FW_IN") 

                    nic_rules << filter_established(chain, :accept)

                    rules.split(" ").each do |rule|

                        nic_rules << filter_rule(chain, rule)

                    end

                    nic_rules << filter_all(chain, :drop)

                end

                if rules = nic[:fw_out]

                    chain = "one-#{vm_id}-#{ip}-out"

                    nic_rules << new_chain(chain)

                    nic_rules << tap_to_chain(tap, chain, "FW_OUT")

                    nic_rules << filter_established(chain, :accept)

                    rules.split(" ").each do |rule|

                        nic_rules << filter_rule(chain, rule)

                    end

                    nic_rules << filter_all(chain, :drop)

                end

                iptables_out = `#{COMMANDS[:iptables]} -L FORWARD`

                if ! iptables_out.match(/^FW_IN/) || ! iptables_out.match(/^FW_OUT/)

                    `#{COMMANDS[:iptables]} -N FW_IN >/dev/null 2>&1`

                    `#{COMMANDS[:iptables]} -N FW_OUT >/dev/null 2>&1`

                    `#{COMMANDS[:iptables]} -A FORWARD -j FW_IN >/dev/null 2>&1`

                    `#{COMMANDS[:iptables]} -A FORWARD -j FW_OUT >/dev/null 2>&1`

                end

                run_rules(nic_rules)

            end

        end

        unlock

    end

    def deactivate

        lock

        vm_id =  @vm['ID']

        process do |nic|

            ip = nic[:ip]

            purge_chain("one-#{vm_id}-#{ip}-in", "FW_IN")

            purge_chain("one-#{vm_id}-#{ip}-out", "FW_OUT")

        end

        unlock

    end

    def purge_chain(chain, parent)

        rules = Array.new

        rules << rule(`#{COMMANDS[:iptables]}-save | grep #{parent} | grep #{chain}`.gsub!(/^-A/,"-D"))

        rules << rule("-F #{chain}")

        rules << rule("-X #{chain}")

        run_rules(rules)

    end

    def filter_established(chain, policy)

        policy = policy.to_s.upcase

        rule "-A #{chain} -m conntrack --ctstate ESTABLISHED,RELATED -j #{policy}"

    end

    def filter_rule(chain, rule)

        r = rule.split(/,\s*(?=[^\[\]]*(?:\[|$))/)

        p = get_protocol(r[0])

        s = get_src(r[1])

        d = get_dst(r[2])

        a = get_action(r[3])

        rule "-A #{chain} #{p} #{s} #{d} #{a}"

    end

    def filter_all(chain, policy)

        policy = policy.to_s.upcase

        rule "-A #{chain} -j #{policy}"

    end

    def get_protocol (protocol)

        "-p #{protocol}"

    end

    def get_src (src)

        host = get_host("s",src)

        port = get_port("s",src)

        "#{host} #{port}"

    end

    def get_dst (dst)

        host = get_host("d",dst)

        port = get_port("d",dst)

        "#{host} #{port}"

    end

    def get_host (tag, data)

        if host = data.split(":")[0].to_s

            host.gsub!(/\[(.*)\]/,'\1')

            if ! host.eql?("")

                "-#{tag} #{host}"

            end

        end

    end

    def get_port (tag, data)

        if port = data.split(":",2)[1].to_s

            port.gsub!(/\[(.*)\]/,'\1')

            if ! port.eql?("")

                "-m multiport --#{tag}ports #{port}"

            end

        end

    end

    def get_action (action)

        case action

            when "ACCEPT"

                "-j RETURN"

            when "DROP"

                "-j DROP"

        end

    end

    def run_rules(rules)

        rules.flatten.each do |rule|

            OpenNebula.exec_and_log(rule)

        end

    end

    def tap_to_chain(tap, chain, parent)

        case parent

            when "FW_IN"

                rule "-A #{parent} -m physdev --physdev-is-bridged --physdev-out #{tap} -j #{chain}"

            when "FW_OUT"

                rule "-A #{parent} -m physdev --physdev-in #{tap} -j #{chain}"

        end

    end

    def new_chain(chain)

        rule "-N #{chain}"

    end

    def rule(rule)

        "#{COMMANDS[:iptables]} #{rule}".gsub(/[ ]+/,' ')

    end

end