#!/usr/bin/ruby
#
# Details: Fgen subsystem, IPtables rule generator.
# Authors: Dmitry Maksyoma <dmaks@esphion.com>
# Started: 01/04/2008
#
# (c) 2008 Esphion Limited
# Copyright in this document, whether in written, electronic or other
# format including any source code or other computer code set forth
# therein or attached thereto and copyright in all parts thereof is owned
# by Esphion Limited, who reserves all rights therein.
# In particular, no part of this document/computer code may be reproduced,
# copied, stored in a retrieval system, used or transmitted by any means
# whatsoever without the prior written consent of Esphion Limited.
#
# All enquiries in relation thereto should be directed to:
#
# Esphion Ltd
# P.O. Box 300496,
# Albany,
# Auckland,
# New Zealand.
#
# Phone: +64 9 415 0227
# Fax:	 +64 9 415 0228
# Email: info@esphion.com

$:.insert(-2, File.join(File.dirname(File.expand_path($0)), "../lib"))
require 'pattern/entities'
require 'fgen/common'
include Entities

TCP_FLAGS_TO_CHECK = [TCPHeader::SYN_FLAG, TCPHeader::ACK_FLAG, TCPHeader::FIN_FLAG, 
  TCPHeader::RST_FLAG, TCPHeader::URG_FLAG, TCPHeader::PSH_FLAG]

def format_ip(bytes)
  return if bytes.index(nil)

  bytes.join('.')
end

# Split layer to continuous chunks. Return array of
# [offset, [byte1, byte2, ...]] arrays.
def splitLayer layer, baseOffset=0
  bytes, result, startOffset = {}, [], nil
  return result if layer.empty?

  0.upto(layer.keys.max) { |offset|
    unless b = layer[offset]
      startOffset = nil
      next
    end

    startOffset ||= offset
    (bytes[startOffset + baseOffset] ||= []).push b
  }
  bytes.keys.sort.each { |offset|
    result.push [offset, bytes[offset]]
  }
  result
end

def format_pattern(pattern, name, src_ips, dst_ips)
  link = pattern.link
  ipl = pattern.ip
  ippl = pattern.ipp
  payload = pattern.packetPayload
  iplUsedFields, ipplUsedFields = [], []
  rule = []

  if ipl && (proto = ipl.proto)
    proto = IPHeader.proto_desc[proto] || proto
    rule.push "-p #{proto}"
    iplUsedFields.push :proto
  end

  if link && macBytes = link[:srcmac]
    macBytes = macBytes.map { |b|
      break false if b.nil?
      "%02X" % b
    }
    rule.push "-m mac --mac #{macBytes.join(':')}" if macBytes
  end

  if ipl
    # Whole IPs within the pattern take precedence of whatever is passed in
    # src_ips/dst_ips.
    src_ips = [$_] if $_ = format_ip(ipl[:srcip])
    dst_ips = [$_] if $_ = format_ip(ipl[:dstip])
    iplUsedFields.push :srcip if src_ips
    iplUsedFields.push :dstip if dst_ips

    # If src_ips or dst_ips contain no more than one IP, source/destination IP
    # will fit into a single chain.
    if (src_ips.nil? || src_ips.size == 1) && (dst_ips.nil? || dst_ips.size == 1) 
      singleChain = true
      [[:srcip, src_ips, '-s'], [:dstip, dst_ips, '-d']].each { |fld, ar, cliOption|
	ip = ar.first if ar && ar.size == 1
	rule.push "#{cliOption} #{ip}" if ip
      }
    end

    if tos = ipl.tos
      rule.push "-m tos --tos #{tos}"
      iplUsedFields.push :tos
    end

    if totalLen = ipl.total_len
      rule.push "-m length --len #{totalLen}"
      iplUsedFields.push :total_len
    end

    # Note: -f (fragmented packets) option doesn't differ between having and
    # not having IP flags. For example, if the DF (`Don't fragment') flag is
    # set even if we specify `! -f' option, it will match packets with both
    # DF set and unset.

    if ttl = ipl.ttl
      rule.push "-m ttl --ttl-eq #{ttl}"
      iplUsedFields.push :ttl
    end

    # Protocol specific.
    if ippl
      if ipl.proto_icmp && (itype = ippl.itype)
	rule.push "--icmp-type #{itype}"
	ipplUsedFields.push :itype
      elsif ipl.proto_tcp || ipl.proto_udp
	[[:srcport, '--sport'],
	    [:dstport, '--dport']].each { |fld, cliOption|
	  if port = ippl.send(fld)
	    rule.push "#{cliOption} #{port}"
	    ipplUsedFields.push fld
	  end
	}

	if ipl.proto_tcp && (flags = ippl.flags)
	  comp = TCP_FLAGS_TO_CHECK.find_all { |flag| flags & flag > 0 }. \
	    map { |flag| TCPHeader.flag_desc[flag].upcase }.join(',')
	  comp = 'NONE' if comp.empty?

	  rule.push "--tcp-flags ALL #{comp}"
	  ipplUsedFields.push :flags
	end
      end
    end

    # At this point, all possible human-readable options were used. All bytes
    # that are left in IP layer, IP protocol layer and payload layer will be
    # added as raw bytes using `-m string --hex-string' (-m string starts
    # looking for --hex-string since IP layer, thus it's not possible to use
    # it to match link layer bytes). Also, according to Dennis, all ethernet
    # level protocol headers except VLAN and MPLS should be ignored (see bug
    # #12077 for example with L2TP).
    
    # Before deleting bytes used by the above options, find out data offsets
    # for IP layer and IP protocol layer.
    proto = ipl.proto
    iplDataOffset = ipl.dataOffset
    if ippl
      ipplDataOffset = ippl.dataOffset
    elsif proto
      ipplDataOffset = Pattern::IP_PROTO_HEADER_MAP[proto].dataOffset
    end

    # Delete all bytes that were used in human-readable options.
    [[iplUsedFields, ipl], [ipplUsedFields, ippl]].
	each { |flds, hdr|
      next unless hdr
      flds.each { |fld| hdr.delFldBytes fld }
    }

    offsets = []
    offsets.concat splitLayer(ipl.layer)

    if iplDataOffset
      if proto == IPHeader::PROTO_ICMP || proto == IPHeader::PROTO_TCP ||
	  proto == IPHeader::PROTO_UDP
	offsets.concat splitLayer(ippl.layer, iplDataOffset) if ippl

	if payload && ipplDataOffset
	  offsets.concat splitLayer(payload.layer, iplDataOffset + ipplDataOffset)
	end
      else
	# Protocol is not one of ICMP, TCP, UDP; data after IP header goes
	# into payload.
	offsets.concat splitLayer(payload.layer, iplDataOffset) if payload
      end
    end

    # Add to the rule `-m string --hex-string' entries.
    # NOTE: --hex-string only matches successfully with offsets that are
    # less than packet length (as defined by IP header), if --from specifies
    # offset that is >= packet length, the rule will never match. This
    # description is valid for iptables of Linux kernel 2.6.18 (Debian Etch).
    #
    # NOTE: `--from N --to N+1' works just fine, on the other hand `--to N'
    # doesn't. Size of --hex part doesn't affect `--to' option.
    offsets.each { |offset, bytes|
      rule.push "-m string --from #{offset} --to #{offset+1} " \
	"--hex '|#{bytes.map { |b| "%02x" % b }.join}|' --alg kmp"
    }
  end

  # Fit everything into a single chain.
  return "iptables -N #{name}; " \
	 "iptables -A #{name} #{rule.join(' ')} -j DROP" if singleChain

  # Make multiple chains: for source IPs, for destination IPs and one for the
  # rest of conditions.
  chains = [chain = "#{name}-rest"]
  restRule = "iptables -A #{chain} #{rule.join(' ')} -j DROP"
  target = "-g #{chain}"

  if dst_ips
    chains << chain = src_ips ? "#{name}-dst" : name
    dstRule = []
    dst_ips.each { |ip| dstRule << "iptables -A #{chain} -d #{ip} #{target};\n" }
    target = "-g #{chain}"
  end

  if src_ips
    chains << chain = name
    srcRule = [] 
    src_ips.each { |ip| srcRule << "iptables -A #{chain} -s #{ip} #{target};\n" }
  end

  ret = chains.reverse.map { |chain| "iptables -N #{chain};" }.join(' ') << "\n"
  ret << srcRule.join if srcRule
  ret << dstRule.join if dstRule
  ret << restRule
  ret
end

do_format
