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]

            if tap

                init_and_clean(tap)

                nic_rules = Array.new

                if rules = nic[:fw_out]

                    chain = "one-#{vm_id}-#{tap}-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

                if rules = nic[:fw_in]

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

                    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

                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}-.*-out", "FW_OUT")

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

        end

        unlock

    end

    def init_and_clean(tap)

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

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

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

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

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

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

        end

        purge_chain("one-.*-#{tap}-out", "FW_OUT")

        purge_chain("one-.*-#{tap}-in", "FW_IN")

    end

    def purge_chain(regex, parent)

        iptables_out = `#{COMMANDS[:iptables]}-save | grep #{parent} | grep \'#{regex}\'`

        if iptables_out.length != 0

            chain = iptables_out.split(" ").last 

            rules = Array.new

            rules << iptables_out.gsub!(/^-A/,"-D")

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

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

            run_rules(rules)

        end

    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_OUT"

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

            when "FW_IN"

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

        end

    end

    def new_chain(chain)

        rule "-N #{chain}"

    end

    def rule(rule)

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

    end

end