#!/usr/bin/env ruby
# Encoding: ASCII
# ^^^^^^^^^^^^^^^ the encoding is set to ascii so the IPAddr extention tests work :)
#
# Simple BGP talker implemented from:
# * RFC1771
# * RFC1997
# * RFC2385
# * RFC3392
# * RFC4760
#
# See https://github.com/BytemarkHosting/bgpfeeder for later versions and more
# information.
#
# (c) Bytemark Hosting 2009-2016
#
# <matthew@bytemark.co.uk>
#


require 'ipaddr'
require 'socket'
require 'logger'
require 'timeout'
require 'optparse'

class String
  def careful_pop(n)
    raise ArgumentError.new("Want #{n} characters, have only got #{length}") if length < n
    slice!(0,n)
  end
  def pop_char; careful_pop(1).unpack("C")[0]; end
  def pop_net_short; careful_pop(2).unpack("n")[0]; end
  def pop_net_int; careful_pop(4).unpack("N")[0]; end

  # Ruby 1.8 returns an int for string[index] and this code assumes that.
  # This method achieves the same thing in a way that works with all rubies
  def byte_at(idx)
    self[idx..idx].unpack('c').first
  end
end

class IPAddr
  # A handy lookup table so that we can quickly map netmasks to prefix
  # lengths, for example:
  #   0xffffff00 => 24
  # or
  #   0xfffffffffffffffffffffffffffff000 => 116
  #
  MASK_REVERSE = {}
  [32, 128].each do |bits|
    (0..bits).each { |b| MASK_REVERSE[(1 << bits) - (1 << b)] = bits-b }
  end

  raise "Bug in MASK_REVERSE" if
    MASK_REVERSE[0xffffff00] != 24 ||
    MASK_REVERSE[0xfffffffffffffffffffffffffffff000] != 116

  # BGP update messages contain IP prefixes encoded as <length,prefix>
  # tuples where length is 0-4 and prefix is the number of relevant
  # bytes of the IP address, e.g.
  #
  #    10.0.0.0/24 => [1,10]
  #    192.168.0.1/32 => [4,192,168,0,1]
  #    0.0.0.0/0 => [0]
  # or
  #    1111:2222::/32 => [4,17,17,34,34]
  #    ::/0 => [0]
  #
  def bgp_encoded_prefix
    bits = MASK_REVERSE[@mask_addr]
    max_bits = ipv4? ? 32 : 128
    bits.chr + if bits > 0
                 to_i_packed[0..((bits+7)/8)-1]
               else
                 ""
               end

  end

  def hash
    to_i | (@mask_addr << 128)
  end

  # Converts the IP address to a apcke
  def to_i_packed
    if ipv4?
      [to_i].pack("N")
    elsif ipv6?
      [((to_i & (0xffffffff << 96))>>96)].pack("N") +
      [((to_i & (0xffffffff << 64))>>64)].pack("N") +
      [((to_i & (0xffffffff << 32))>>32)].pack("N") +
      [((to_i &  0xffffffff       )    )].pack("N")
    else
      raise ArgumentError.new("Unknown IP address type")
    end
  end
end

raise "Bug in IPAddr extensions" if
  IPAddr.new("1.2.3.4/24").bgp_encoded_prefix != "\030\001\002\003" ||
  IPAddr.new("0.0.0.0/0").bgp_encoded_prefix != "\000" ||
  IPAddr.new("255.255.255.255/32").bgp_encoded_prefix != " \377\377\377\377" ||
  IPAddr.new("::/0").bgp_encoded_prefix != "\000" ||
  IPAddr.new("1111:2222::/32").bgp_encoded_prefix != " \021\021\"\"" ||
  IPAddr.new("1111:2222::/33").bgp_encoded_prefix != "!\021\021\"\"\000"

# Parses IPv4 and IPv6 address / port combinations, allowing the user to specify
# just an IP address, an IP address and port, or two IP addresses and ports
# (the first assumed to be a source address).
#
# e.g. these are valid:
#
#    1.2.3.4              => [IPAddr.new("1.2.3.4"), nil]
#    1.2.3.4:80           => [IPAddr.new("1.2.3.4"), 80]
#    :80                  => [nil, 80]
#    1111:aaaa::1         => [IPAddr.new("1111:aaaa::1"), nil]
#    [1111:AAAA::1]:80    => [IPAddr.new("1111:aaaa::1"), 80]
#    []:80                => [nil, 80]
#
# You can also specify two endpoints:
#
#    127.0.0.1:80-127.0.0.2:4455
#    1111:aaaa::1-[2222:bbbb::2]:80
#
# and your IP
#
# These will throw an ArgumentError:
#
#   foo.bar.com           # no name resolution allowed
#   1111111111            # IPAddr explodes
#   bad string            # not even close
#   1.2.3.4-5.6.7.8-9.10.11.12 # too many dashes
#
class IpAddrAndPort
  attr_reader :ipaddr, :port
  attr_reader :from

  # Parse an IP address/port combination, optionally specifying a default port.
  #
  def initialize(spec, default_port=nil)
    raise ArgumentError.new("default_port must be a number") unless
     default_port.nil? || default_port.kind_of?(Fixnum)

    spec1, spec2 = spec.split(/-/, 2)
    @from=nil

    if spec2
      @from = IpAddrAndPort.new(spec1, nil)
      spec = spec2
    end

    @ipaddr, @port = parse(spec)
    @port ||= default_port

    if @from
      raise ArgumentError.new("Only one dash allowed in IP spec (left-hand side means 'source', right-hand side means 'destination'") if @from.from
      raise ArgumentError.new("Can't connect an IPv4 source to an IPv6 destination (or vice versa)") if
        @from.ipaddr.ipv6? != ipaddr.ipv6?
    end
  end

  # Returns a two-element array of [IPAddr, port]
  #
  def to_a
    [ipaddr, port]
  end

  # Returns the IP address and port as a string.
  #
  def to_s
    dest_s = if ipaddr
      if ipaddr.ipv6?
        if port
          "[#{ipaddr}]:#{port}"
        else
          ipaddr.to_s
        end
      else
        ipaddr.to_s + (port ? ":#{port}" : "")
      end
    else
      ":#{port}"
    end

    if from
      from.to_s + "-" +dest_s
    else
      dest_s
    end
  end

  def parse(spec)
    spec = spec.downcase.gsub(/^\s+|\s+$/, "")
    port = nil

    if /^(.*\]):(\d+)$/.match(spec) ||
       /^([\d\.]*):(\d+)$/.match(spec)
      port = $2.to_i
      spec = $1
    end
    if spec =~ /^\[([^\]]+)\]$/
      spec = $1
    end

    if spec.empty?
      ipaddr = nil
    else
      if !/^\d+\.\d+\.\d+\.\d+$/.match(spec) && #ipv4
         !/^[a-z0-9\:]+$/.match(spec) # ipv6, roughly
        raise ArgumentError.new("#{spec} is not an IPv4 or IPv6 address")
      else
        begin
          ipaddr = IPAddr.new(spec)
        rescue ArgumentError => ae
          raise ArgumentError.new("Couldn't parse IP #{spec}")
        end
      end
    end

    [ipaddr, port]
  end
end

# A Linux-specific class to set up a RFC2385-compliant MD5-secured TCP
# connection.
#
class TCPMD5Socket < Socket
  IPPROTO_TCP =  6 # linux/in.h
  TCP_MD5SIG  = 14 # linux/tcp.h
  TCP_MD5SIG_MAXKEYLEN = 80 # linux/tcp.h

  # Works exactly the same as TCPSocket.open except you can supply a password
  # to pass to the kernel for MD5 authhentication.
  #
  def initialize(password, host, port, local_host=nil, local_port=nil)

    raise ArgumentError.new("Password is too long") if
      password.length > TCP_MD5SIG_MAXKEYLEN

    family = Socket.const_get(IPAddr.new(host).ipv4? ? "AF_INET" : "AF_INET6")
    super(family, Socket::SOCK_STREAM, 0)
    # struct tcp_md5sig {
    # 	struct __kernel_sockaddr_storage tcpm_addr;	/* address associated */
    # 	__u16	__tcpm_pad1;				/* zero */
    # 	__u16	tcpm_keylen;				/* key length */
    # 	__u32	__tcpm_pad2;				/* zero */
    # 	__u8	tcpm_key[TCP_MD5SIG_MAXKEYLEN];		/* key (binary) */
    # };
    tcp_md5sig_buffer = [
      Socket.pack_sockaddr_in(port, host), 0, password.length, 0, password
    ].pack("a128SSLa#{TCP_MD5SIG_MAXKEYLEN}")
    setsockopt(IPPROTO_TCP, TCP_MD5SIG, tcp_md5sig_buffer)
    bind(Socket.pack_sockaddr_in(local_port, local_host)) if local_host || local_port
    connect(Socket.pack_sockaddr_in(port, host))
  end
end

module BGP
  REGEXP_IP_ADDRESS = /^\d+\.\d+\.\d+\.\d+$/
  REGEXP_IP_MASK    = /^\d+\.\d+\.\d+\.\d+\/(?:\d+|\d+\.\d+\.\d+\.\d+)$/

  # Must set BGP::Log externally, oops, maybe that's ugly

  module Notification
    CODES = [
             nil,
             "Message Header Error",
             "OPEN Message Error",
             "UPDATE Message Error",
             "Hold Timer Expired",
             "Finite State Machine Error",
             "Cease"
            ]

    SUBCODES = [
                [],
                [nil, # Message Header Error subcodes
                 "Connection Not Synchronized",
                 "Bad Message Length",
                 "Bad Message Type",
                ],
                [nil, # OPEN Message Error subcodes
                 "Unsupported Version Number",
                 "Bad Peer AS",
                 "Bad BGP Identifier",
                 "Unsupported Optional Parameter",
                 "Authentication Failure",
                 "Unacceptable Hold Time",
                 "Unsupported Capability",
                 "Grouping Conflict",
                 "Grouping Required",
                 "Redirecting Now",
                 "Redirect Required"
                ],
                [nil, #UPDATE Message Error subcodes
                 "Malformed Attribute List",
                 "Unrecognized Well-known Attribute",
                 "Missing Well-known Attribute",
                 "Attribute Flags Error",
                 "Attribute Length Error",
                 "Invalid ORIGIN Attribute",
                 "AS Routing Loop",
                 "Invalid NEXT_HOP Attribute",
                 "Optional Attribute Error",
                 "Invalid Network Field",
                 "Malformed AS_PATH"
                ]
               ]

    class General < Exception
      attr_reader :code, :subcode

      def sent?; @sent; end
      def received?; !sent?; end

      def initialize(code, subcode, sent=false)
        @code = code
        @subcode = subcode
        @sent = !!(sent)
        code_msg, subcode_msg = General.lookup_messages(code, subcode)
        subcode_msg ||= "subcode #{subcode}"
        super((sent? ? "sent " : "received ") + "#{code_msg} (#{subcode_msg})")
      end

      def self.lookup_messages(code, subcode)
        code_msg = CODES[code]
        subcode_msg = nil
        if CODES[code] && SUBCODES[code] && subcode
          subcode_msg = SUBCODES[code][subcode]
        end
        [code_msg, subcode_msg]
      end

      def self.lookup_class_name(code, subcode)
        code_msg, subcode_msg = General.lookup_messages(code, subcode)
        name = subcode_msg ? subcode_msg : code_msg
        name.split(/[^a-zA-Z]/).map { |word| word.capitalize }.join
      end
    end

    def self.define_new_exception(code_or_parent, subcode=nil)
      parent = code_or_parent.kind_of?(Class) ? code_or_parent : General
      code = code_or_parent.kind_of?(Class) ? code_or_parent.new(1).code : code_or_parent
      name = General.lookup_class_name(code, subcode)
      raise "Unown BGP Exception #{code}.#{subcode}" if !name
      if parent == General
        class_eval "class #{name} < #{parent}; def initialize(subcode,sent=false); super(#{code},subcode,sent); end; end"
      else
        class_eval "class #{name} < #{parent}; def initialize(sent=false); super(#{subcode},sent); end; end"
      end
      const_get(name.to_sym)
    end

    def self.raise_by_code(code, subcode, sent)
      clazz = const_get(General.lookup_class_name(code, subcode))
      if clazz.superclass == General
        raise clazz.new(subcode, sent)
      else
        raise clazz.new(sent)
      end
    end

    CODES.each_with_index do |name1, code|
      next unless name1
      parent = define_new_exception(code)
      next unless SUBCODES[code]
      SUBCODES[code].each_with_index do |name2, subcode|
        next unless name2
        define_new_exception(parent, subcode)
      end
    end
  end

  raise "Error defining exceptions" unless
    BGP::Notification::General.new(1,1,true).code == 1 &&
    BGP::Notification::General.new(1,1,true).subcode == 1 &&
    BGP::Notification::General.new(1,1,false).received? &&
    BGP::Notification.constants.include?(:UpdateMessageError) &&
    BGP::Notification.constants.include?(:InvalidOriginAttribute) &&
    !BGP::Notification.constants.include?(:Subcode) &&
    BGP::Notification::UpdateMessageError.new(1,true).sent? &&
    BGP::Notification::UpdateMessageError.new(1,false).received? &&
    BGP::Notification::UpdateMessageError.new(99).code == 3 &&
    BGP::Notification::UpdateMessageError.new(99).subcode == 99 &&
    BGP::Notification::InvalidOriginAttribute.new.code == 3 &&
    BGP::Notification::InvalidOriginAttribute.new.subcode == 6 &&
    BGP::Notification::InvalidOriginAttribute.superclass == BGP::Notification::UpdateMessageError

  # Simple data structure for representing a BGP route update, or at
  # least all the features we support.  The expectation is that
  # prefix, next_hop and withdraw may represent IPv4 or IPV6
  # addresses, and the BGP-speaking code below is expected to separate
  # this into the MP_ attributes as appropriate.
  #
  # http://justrudd.org/2007/05/02/ruby-weirdness/ bit me when trying to do
  # array differences, so need to have hash & eql? implemented here and in
  # the IPAddr class (see above also).
  #
  class Update < Struct.new(:prefix, :next_hop, :withdraw, :local_pref, :communities)
    def initialize(prefix, next_hop=nil)
      self.communities = []
      self.prefix = prefix
      self.next_hop = next_hop
    end
    def ipv6?
      prefix.ipv6?
    end
    def withdraw?
      next_hop.nil? || withdraw
    end
    def new_route?
      next_hop && withdraw.nil?
    end
    def replace_route?
      next_hop && withdraw
    end
    def prefix=(p)
      p = IPAddr.new(p) if p && !p.kind_of?(IPAddr)
      super(p)
    end
    def next_hop=(h)
      h = IPAddr.new(h) if h && !h.kind_of?(IPAddr)
      super(h)
    end
    def communities=(list)
      super(Update.parse_communities(list))
    end

    # Help with parsing textual descriptions
    #
    _IP = '(?:\d+\.\d+\.\d+\.\d+|[A-Fa-f0-9:]+)'
    FORMAT_LINUX = /^(#{_IP}\/\d+)\s+(?:via )(#{_IP})/
    FORMAT_CISCO = /^(?:ip route )?(#{_IP}) (#{_IP}) (#{_IP})/
    FORMAT_COMMUNITIES = /\Wcommunities\s+([\d\:\,]+)/
    FORMAT_LOCAL_PREF  = /\Wlocal_pref\s+(\d+)/

    def self.parse_line(line)
      case line
        when FORMAT_LINUX then base=Update.new($1,$2)
        when FORMAT_CISCO then base=Update.new("#{$1}/#{$2}", $3)
        else
          return nil
      end
      base.communities = $1.split(",") if FORMAT_COMMUNITIES.match(line)
      base.local_pref = $1.to_i if FORMAT_LOCAL_PREF.match(line)
      base
    end

    def self.parse_communities(list)
      list = list.split(/[ ,]+/) if list.kind_of?(String)
      raise ArgumentError.new("Must pass String or Array") unless
        list.kind_of?(Array)

      list.map do |community|
        case community
          when Integer then community
          when /^0x([0-9a-fA-F]{1,8})$/ then community.hex
          when /^\d+$/ then community.to_i
          when /^(\d+):(\d+)$/ then ($1.to_i << 16) | $2.to_i
          else
            raise ArgumentError.new("Bad community spec '#{community}'")
        end
      end
    end

    def equal_except_withdraw?(b)
      a1 = entries
      b1 = b.entries
      a1[2] = b1[2] = nil
      a1 == b1
    end
    def eql?(b)
      entries == b.entries
    end
    def hash
      (prefix.hash | next_hop.hash << 32) ^ (withdraw ? 0xffffffff : 0 )
    end
  end

  class Peer
    attr_reader :remote

    # Initialise a new peer based on the given Database, and peer_spec.  The
    # peer_spec can be one of:
    #
    #   "1.2.3.4"
    #   "1.2.3.4:1799"
    #   "1.2.3.4 my_session_secret"
    #
    # i.e. it is a list of IP addresses, but you can also specify:
    # * a TCP port number (by appending e.g. :1799 to the IP);
    # * a TCP-MD5 password (as a second argument).
    #
    #
    def initialize(database, peer_spec, initial_sleep=nil)
      @database = database
      @hold_time = database.hold_time # possibly overridden by remote OPEN

      peer_spec, @tcp_md5_secret = peer_spec.split(/\s+/,2)
      @remote = IpAddrAndPort.new(peer_spec, 179)

      @initial_sleep = initial_sleep

      @closed = false
    end

    def close
      @closed = true
    end

    def run
      begin
        if @initial_sleep
          Log.warn("Sleeping for #{@initial_sleep}s due to previous error")
          sleep @initial_sleep
        end
        real_run
      rescue Exception => ex
        Log.error("Uncaught exception in peer thread, re-raising: #{ex}")
        Log.error(ex.backtrace.join("\n"))
        raise
      end
    end

    def real_run
      while !@closed
        started = Time.now
        connect_error_reported = false
        advertise_capabilities = true

        begin
          Timeout::timeout(10) do
            @s = if @tcp_md5_secret
                   if remote.from
                     TCPMD5Socket.new(@tcp_md5_secret, remote.ipaddr.to_s, remote.port,
                                      remote.from.ipaddr.to_s, remote.from.port)
                   else
                     TCPMD5Socket.new(@tcp_md5_secret, remote.ipaddr.to_s, remote.port)
                   end
                 else
                   if remote.from
                     TCPSocket.new(remote.ipaddr.to_s, remote.port,
                                   remote.from.ipaddr.to_s, remote.from.port)
                   else
                     TCPSocket.new(remote.ipaddr.to_s, remote.port)
                   end
                 end
          end
          if connect_error_reported
            l_warn("Connection established after #{Time.now - started}s")
          end

          # returns after OPEN messages have been received both ways
          do_open(advertise_capabilities)

        rescue Notification::General => notification
          @s.close
          @s = nil

          if advertise_capabilities && (notification.kind_of?(Notification::OptionalAttributeError) ||
                                        notification.kind_of?(Notification::UnsupportedCapability) ||
                                        notification.kind_of?(Notification::GroupingConflict)
                                        )
            l_warn("Peer doesn't do multiprotocol, trying again without")
            advertise_capabilities = false
            retry
          end

          l_warn("BGPv4 connection failed during OPEN, will try again in 5 mins: #{notification}") unless connect_error_reported
          connect_error_Reported = true
          sleep 300
          break if @closed
          retry

        # Try to predict some errors that we'll see if BGP peers come and go
        # on the network - others will abort the peer thread with messier
        # error.
        #
        rescue Timeout::Error,
               Errno::ECONNREFUSED,
               Errno::ETIMEDOUT,
               Errno::EHOSTDOWN,
               Errno::EHOSTUNREACH,
               Errno::ENETDOWN,
               Errno::ENETUNREACH,
               Errno::ENETRESET,
               Errno::ECONNABORTED,
               Errno::ECONNRESET => e

          @s = nil
          if !connect_error_reported
            l_warn("Failed BGPv4 connection due to #{e}, will keep trying every 10s")
            connect_error_reported = true
          end
          sleep 10
          break if @closed
          retry

        end

        if @s && !@closed
          begin
            # returns when connection is ready to be closed by request
            do_maintain_connection
          rescue Notification::General => notify
            l_warn("BGPv4 connection terminated: #{notify}")
          end
          @s.close
        end

        if !@closed
          retry_delay = 20 - (Time.now - started)
          retry_delay = 0 if retry_delay < 1
          l_warn("died, retrying "+
            (retry_delay == 0 ? "immediately" : "in #{retry_delay}s")
          )
          sleep retry_delay
        end
      end
    end

    protected

    def l_debug(msg); Log.debug("peer #{remote}: #{msg}"); end
    def l_error(msg); Log.error("peer #{remote}: #{msg}"); end
    def l_warn(msg); Log.warn("peer #{remote}: #{msg}"); end
    def l_info(msg); Log.info("peer #{remote}: #{msg}"); end

    TYPE_OPEN = 1
    TYPE_UPDATE = 2
    TYPE_NOTIFICATION = 3
    TYPE_KEEPALIVE = 4

    # first communication once socket is open, send OPEN, expect to
    # receive the same from peer, note its hold time.
    #
    def do_open(advertise_capabilities=false)
      send_open(advertise_capabilities)
      type, message = read_next_message
      if type == TYPE_NOTIFICATION
        handle_notification(message)
        false
      else
        send_notification(1,1,"Expected OPEN, got #{type} instead") if type != TYPE_OPEN
        send_notification(2,1,"Unsupported BGP version, peer is speaking v#{message.byte_at(0)}") if message.byte_at(0) != 4
        @peer_autonomous_system = message[1..2].unpack("n")[0]
        l_info "Peer is AS#{@peer_autonomous_system}, assuming "+
          (ebgp? ? "exterior" : "interior")+
          " BGP conventions"
        peer_hold_time = message[3..4].unpack("n")[0]
        peer_as = message[5..8].unpack("N")[0]
        @hold_time = peer_hold_time if peer_hold_time < @database.hold_time
        l_debug "hold time for connection is #{@hold_time}s, sending first KEEPALIVE"
        handle_optional_parameters(message[9..-1])
        send_keepalive
        true
      end
    end

    # The only optional parameter we support is Capabilities Advertisement, we
    # object to anything else.
    #
    def handle_optional_parameters(raw)
      @peer_ipv6_capable = false

      send_notification(1,1,"Optional parameters missing") if raw.empty?
      length = raw.pop_char
      send_notification(1,1,"Optional parameters length declared #{length}, is #{raw.length}") unless
        length == raw.length

      while !raw.empty? do
        param_type = raw.pop_char
        param_length = raw.pop_char
        param = raw.careful_pop(param_length)

        send_notification(2,4,"Don't know Optional Parameter #{param_type}") unless param_type == 2

        while !param.empty?
          cap_code = param.pop_char
          cap_length = param.pop_char
          cap_value = param.careful_pop(cap_length)

          send_notification(1,0,"cap_value is #{cap_value.length}, declared to be #{cap_length}") if
            cap_length != cap_value.length

          if cap_code == 1 # Multiprotocol Extensions RFC 2858
            send_notification(1,0,"multiprotocol capability size should be 4") unless
              cap_value.length == 4
            afi, _reserved, safi = cap_value.unpack("ncc")

            if afi == 2 && safi == 1
              l_info("Peer understands IPv6/unicast")
              @peer_ipv6_capable = true
            else
              l_warn("Unknown multiprotocol capability ignored afi=#{afi}, safi=#{safi}")
            end
          else
            l_warn("Unknown capability #{cap_code} ignored")
          end
        end
      end
    end

    def ibgp?
      @peer_autonomous_system == @database.autonomous_system
    end
    def ebgp?; !ibgp?; end

    # Simple polling loop to check for updates and respond to incoming
    # messages appropriately.  Can probably fix this to use IO.select
    # but maybe not necessary.
    #
    def do_maintain_connection
      warned_of_received_updates = false
      last_database_update = nil

      while !@closed
        # 1: send keepalive if it's due
        next_keepalive_at = @alive_last_sent + (@hold_time/3)
        send_keepalive if next_keepalive_at <= Time.now

        # 2: send updates if they're available (first call should give us
        # all entries)
        updates = @database.updates_since(last_database_update)
        if updates.length > 0
          send_updates(updates)
          last_database_update = Time.now
        end

        # 3: handle incoming messages, waiting maximum 1s
        type, message = read_next_message(1)

        case type
          when TYPE_OPEN
            send_notification(5,0,"Received second OPEN message")

          when TYPE_UPDATE
            if !warned_of_received_updates
              l_warn("ignoring #{message.length} byte UPDATE from peer (further UPDATES will not be warned about, fix peer configuration!)")
              warned_of_received_updates = true
            end

          when TYPE_NOTIFICATION
            handle_notification(message)

          when TYPE_KEEPALIVE
            # @alive_last_seen set automatically by read_next_message

          when nil
            # timeout, nothing received, go round again

          else
            send_notification(1,3,"Unknown message type #{type} received")
        end

        # 4: notify and terminate connection if hold timer has expired
        send_notification(4,0,"Hold timer expired") if
          Time.now - @alive_last_seen > @hold_time
      end

      send_notification(6,0,"User closed connection, peer will finish",false)
    end

    def handle_notification(message)
      code, subcode, data = message.byte_at(0), message.byte_at(1), message[2..7]
      Notification.raise_by_code(code, subcode, false)
    end

    def send_open(advertise_capabilities=false)
      l_debug "sending OPEN for AS#{@database.autonomous_system} with hold time #{@database.hold_time} and BGP ID #{@database.bgp_identifier}"

      open_data = [
              4, # version, BGP 4
              @database.autonomous_system,
              @database.hold_time,
              @database.bgp_identifier,
             ].pack("cnnN")

      optional_data = ""
      if advertise_capabilities
        capability_data = [
                           1, # capability code 1 (RFC2858)
                           4, # capability length
                           2, # AFI: protocol code for IPv6
                           0, # reserved
                           1  # SAFI: indicate unicast
                           ].pack("ccncc")

        optional_data = 2.chr +  # parameter type 2 (RFC3392)
          capability_data.length.chr +
          capability_data
      end

      send_message(TYPE_OPEN, open_data + optional_data.length.chr + optional_data)
    end

    ATTR_ORIGIN = 1
    ATTR_AS_PATH = 2
    ATTR_NEXT_HOP = 3
    ATTR_LOCAL_PREF = 5
    ATTR_COMMUNITIES = 8
    ATTR_MP_REACH_NLRI = 14
    ATTR_MP_UNREACH_NLRI = 15

    AS_SET = 1
    AS_SEQUENCE = 2

    ATYPE_OPTIONAL = 128
    ATYPE_TRANSITIVE = 64
    ATYPE_PARTIAL = 32
    ATYPE_EXTENDED_LENGTH = 16

    def send_updates(updates)
      updates = updates.reject { |u| u.ipv6? } unless @peer_ipv6_capable

      # There is scope to aggregate updates a little bit here, but probably
      # would only clutter the code; as long as caller has ensured that
      # "replaced" routes are sent as one update, we avoid flaps which are
      # the main hazard.
      #
      l_debug "send_updates: #{updates.inspect}"

      updates.each do |update|
        update_data = update.ipv6? ? update_data_multiprotocol(update) :
          update_data_ipv4(update)

        # send the complete message
        send_message(TYPE_UPDATE, update_data)
      end
    end

    def update_data_multiprotocol(update)
      update_data = ""

      # 1: withdrawn routes, always 0 as we use the MP_ attributes instead
      #
      update_data += "\0\0"

      # 2: path attributes - use MP_ attributes instead, but otherwise follows the
      # pattern of update_data_ipv4.
      #
      path_attributes_raw = ""

      # 2a: withdrawn routes first
      path_attributes_raw += attribute_raw_unreach_nlri(update.prefix) if
        update.withdraw?

      # 2b: now supply the new feasible route
      if update.next_hop
        path_attributes_raw += attribute_raw_origin
        path_attributes_raw += attribute_raw_as_sequence

        # only send LOCAL_PREF over interior connections
        path_attributes_raw += attribute_raw_local_pref(update.local_pref) if
          ibgp?

        path_attributes_raw += attribute_raw_communities(update.communities) unless
          update.communities.empty?

        path_attributes_raw += attribute_raw_reach_nlri(update.prefix, update.next_hop)
      end

      update_data += [path_attributes_raw.length].pack("n") + path_attributes_raw

      # 3: NLRI information is always empty as we use the MP_ attribute instead
      #

      update_data
    end

    # Generate unadorned IPv4 BGP updates
    #
    def update_data_ipv4(update)
      update_data = ""
      # 1: withdrawn routes, either 0 or 1 in this implementation
      #
      withdrawn_raw = update.withdraw? ?
      update.prefix.bgp_encoded_prefix :
        ""
      update_data += [withdrawn_raw.length].pack("n") + withdrawn_raw

      # l_debug "withdrawn data: #{withdrawn_raw.inspect}"

      # 2: path attributes
      #
      # The only attribute that provides different information is NEXT_HOP
      # (obviously) but we are obliged to supply ORIGIN (set to internal
      # gateway protocol) and LOCAL_PREF which is set by the database.
      # AS_PATH is zero-length if we're talking to an interior BGP peer,
      # and set to our local AS# if we're talking to an external peer.
      #
      if update.next_hop
        # always send ORIGIN, AS_PATH and NEXT_HOP
        path_attributes_raw =
          attribute_raw_origin +
          attribute_raw_as_sequence +
          attribute_raw_next_hop(update.next_hop)

        # only send LOCAL_PREF over interior connections
        path_attributes_raw += attribute_raw_local_pref(update.local_pref) if
          ibgp?

        path_attributes_raw += attribute_raw_communities(update.communities) unless
          update.communities.empty?

        update_data += [path_attributes_raw.length].pack("n") + path_attributes_raw

        # l_debug "path attributes data: #{path_attributes_raw.inspect}"

        # 3: prefixes that the above path attributes apply to, always 1
        # per update unless we're just withdrawing a previous route.
        #
        update_data += update.prefix.bgp_encoded_prefix
      else
        # zero length path attributes if we're just withdrawing
        update_data += "\0\0"
      end

      # l_debug "prefix data: #{update.prefix.bgp_encoded_prefix.inspect}"
      update_data
    end

    def attribute_raw_as_sequence
      # AS path is empty over IBGP, otherwise just set to [local_as]
      as_path_segments = ibgp? ? "" :
        [AS_SEQUENCE, 1, @database.autonomous_system].pack("ccn")
      [ATYPE_TRANSITIVE, ATTR_AS_PATH, as_path_segments.length].pack("ccc") +
        as_path_segments
    end


    def attribute_raw_origin
      [ATYPE_TRANSITIVE, ATTR_ORIGIN, 1, 0].pack("cccc")
    end

    def attribute_raw_next_hop(addr)
      [ATYPE_TRANSITIVE, ATTR_NEXT_HOP, 4, addr.to_i].pack("cccN")
    end

    def attribute_raw_local_pref(value)
      [ATYPE_TRANSITIVE, ATTR_LOCAL_PREF, 4, value].pack("cccN")
    end

    def attribute_raw_communities(list)
      communities_raw = list.map { |c| [c].pack("N") }.join
      [ATYPE_TRANSITIVE | ATYPE_OPTIONAL, ATTR_COMMUNITIES].pack("cc") +
        communities_raw.length.chr +
        communities_raw
    end

    def attribute_raw_reach_nlri(v6_prefix, v6_next_hop)
      reach_nlri_raw = [2,1].pack("nc") + # AFI,SAFI
        16.chr + # length in BYTES of next hop address
        v6_next_hop.to_i_packed +
        0.chr + #no SNPAs
        v6_prefix.bgp_encoded_prefix

      [ATYPE_OPTIONAL, ATTR_MP_REACH_NLRI, reach_nlri_raw.length].
        pack("ccc") + reach_nlri_raw
    end

    def attribute_raw_unreach_nlri(v6_prefix)
      unreach_nlri_raw = [2,1].pack("nc") + # AFI,SAFI
        v6_prefix.bgp_encoded_prefix
      [ATYPE_OPTIONAL, ATTR_MP_UNREACH_NLRI, unreach_nlri_raw.length].
        pack("ccc") + unreach_nlri_raw
    end

    def send_keepalive
      # l_debug { "sending KEEPALIVE" }
      send_message(TYPE_KEEPALIVE)
      @alive_last_sent = Time.now
    end

    # Notifications are always fatal
    def send_notification(code, subcode, explanation=nil, is_error=true)
      if is_error
        l_error(explanation) if explanation
      else
        l_warn(explanation) if explanation
      end
      send_message(TYPE_NOTIFICATION,
        [code, subcode].pack("cc") + "\0\0\0\0\0\0"
      )
      Notification.raise_by_code(code, subcode, true)
    end

    # Read next BGPv4 message, checking the marker and updating
    # @alive_last_seen so we know not to panic.
    #
    def read_next_message(timeout_seconds=0)
      begin
        Timeout::timeout(timeout_seconds) do
          header = @s.read(19)
          header = header.force_encoding(Encoding::ASCII_8BIT) if header.respond_to?(:force_encoding)
          return nil unless header
          # l_debug { "received header: "+header.inspect }
          send_notification(1,1,"Marker not detected (or out of sync)") if header[0..15] != MARKER
          length, type = header[16..18].unpack("nc")
          if type == TYPE_OPEN || type == TYPE_UPDATE || type == TYPE_KEEPALIVE
            @alive_last_seen = Time.now
          end
          data = @s.read(length-19)
          data = data.force_encoding(Encoding::ASCII_8BIT) if data.respond_to?(:force_encoding)
          # l_debug { "received data: "+data.inspect }
          return [type, data]
        end
      rescue Timeout::Error => timeout
        # l_debug "timeout!"
      end
      nil
    end

    # Send BGPv4 message, adding appropriate header (19 is the length of the
    # marker plus message header which needs to be included)
    def send_message(type, message="")
      data = MARKER + [19 + message.length, type].pack("nc") + message
      # l_debug { "sent: "+data.inspect }
      @s.write(data)
    end

    def next_keepalive_at
      ideal_time = @hold_time / 3
      ideal_time = 1 if ideal_time == 0
      @alive_last_sent + ideal_time
    end

    MARKER = "\377"*16

  end

  # Utility class to watch a file, let us know when it has been updated, and
  # provide a quick way to read it into memory.
  #
  class FileWatcher
    attr_reader :last_updated

    def initialize(file)
      raise Errno::ENOENT unless File.exists?(file)
      @file = file
      @last_updated = nil
    end

    def updated?
      updated_at = File.stat(@file).mtime
      if updated_at != @last_updated
        @last_updated = updated_at
        BGP::Log.debug "#{@file} changed at #{File.stat(@file).mtime}"
        true
      else
        false
      end
    end

    def read_lines
      File.open(@file) do |fh|
        fh.read.split("\n")
      end
    end
  end

  class Database
    attr_reader :autonomous_system
    attr_reader :bgp_identifier

    attr :hold_time, true
    attr :local_pref, true
    attr :communities, true

    def initialize(autonomous_system, bgp_identifier, file)
      @autonomous_system = autonomous_system
      @bgp_identifier = case bgp_identifier
                        when /[\.\:]/ then IPAddr.new(bgp_identifier).to_i & 0xffffffff
                        when /^\d+$/ then bgp_identifier.to_i & 0xffffffff
                        else
                          raise ArgumentError.new("Can't parse BGP identifier #{bgp_identifier}")
                        end
      @hold_time = 60
      @local_pref = 100
      @communities = []
      @file = FileWatcher.new(file)
      @routes = []
      @route_updates = []
    end

    def updates_since(requested_time)
      poll_for_updates
      return @routes if requested_time.nil?

      # inefficient but clear!
      @route_updates.inject([]) do |list, (time_of_update, update)|
        if time_of_update < requested_time
          list
        else
          list + update
        end
      end
    end

    protected

    def poll_for_updates
      return unless @file.updated?
      new_routes = full_update_from_lines(@file.read_lines)

      # add system defaults to each route
      #
      new_routes.each do |route|
        route.local_pref ||= local_pref
        route.communities += communities
      end

      routes_to_add    = new_routes - @routes
      routes_to_delete = @routes - new_routes

      routes_to_change = []
      # try to combine adds & deletes into a single "change" update to avoid
      # causing a flap.  FIXME: maybe a smarter/indexed search would make
      # this faster, can't see it being necessary for our use though?
      #
      routes_to_delete.each do |route|
        routes_to_add.each do |route2|
          if route.equal_except_withdraw?(route2)
            routes_to_add.delete(route2)
            routes_to_delete.delete(route)
            route2.withdraw = true
            routes_to_change << route2
          end
        end
      end

      Log.info "routes file updated at #{@file.last_updated}: "+
        "#{routes_to_add.length} added, "+
        "#{routes_to_delete.length} deleted, "+
        "#{routes_to_change.length} changed" unless
        routes_to_add.empty? &&
        routes_to_delete.empty? &&
        routes_to_change.empty?

      routes_to_delete.each { |route| Log.warn "deleting route #{route}" }
      routes_to_add.each { |route| Log.warn "adding route #{route}" }
      routes_to_change.each { |route| Log.warn "changing route #{route}" }

      # rewrite since we've pulled this out of our current route table
      routes_to_delete.each do |route|
        route.withdraw = true
        route.next_hop = nil
      end

      # Add this calculated set of route updates to our log
      #
      # NB we must always send withdraw requests before we add new routes,
      # in case the prefix is the same but the other attributes differ.
      #
      @route_updates << [@file.last_updated,
        routes_to_delete +
        routes_to_add +
        routes_to_change
      ]

      # finally overwrite the full route table with the newly-parsed one
      @routes = new_routes

      Log.warn "routes file is empty" if @routes.empty?

      nil
    end

    def full_update_from_lines(lines)
      lines.map { |l| Update.parse_line(l) }.compact
    end
  end

  class PeerManager
    def initialize(file, database)
      @file = FileWatcher.new(file)
      @database = database
      @peer_threads = {}
      @close_all = false
    end

    def close_all
      @close_all = true
    end

    def run
      while !@close_all || !@peer_threads.empty?
        if @file.updated?

          # Work out which peers have been added, and which removed.  We
          # treat a peer line with a different TCP-MD5 secret as a different
          # peer, which means the session will be reestablished if the user
          # changes the secret.
          #
          peer_list_new = @file.read_lines.
            select { |l| /^[\[\]\-a-fA-F\d+\.\:]+( .*)?$/.match(l) }
          peer_list_finished = @peer_threads.keys - peer_list_new
          peer_list_to_start = peer_list_new - @peer_threads.keys

          # Kill off the peers we've finished with
          #
          peer_list_finished.each do |peer_spec|
            Log.info "closing peer #{peer_spec}"
            @peer_threads[peer_spec][:peer].close
          end

          # Initialise new ones
          #
          peer_list_to_start.each do |peer_spec|
            Log.info "starting new peer #{peer_spec}"
            peer = Peer.new(@database, peer_spec)
            @peer_threads[peer_spec] = Thread.new { peer.run }
            @peer_threads[peer_spec][:peer] = peer
          end

          Log.warn("peers file is empty, nothing to do for now") if
            peer_list_new.length == 0
        end

        if @close_all
          Log.info("closing all peer connections")
          @peer_threads.each do |name, peer_thread|
            peer_thread[:peer].close
            peer_thread.run
          end
        end

        sleep 1

        # Reap any that have finished
        #
        @peer_threads.each do |peer_spec, thread|
          next if thread.alive?
          begin
            thread.join
            Log.info "peer thread #{peer_spec} finished normally"
            @peer_threads.delete(peer_spec)
          rescue Exception => ex
            Log.error "peer thread #{peer_spec} finished with an exception, will restart in 5 mins but may indicate a bug: #{ex}"
            Log.error ex.backtrace.join("\n")+"\n"
            resurrected_peer = Peer.new(@database, peer_spec, 300)
            @peer_threads[peer_spec] = Thread.new { resurrected_peer.run }
            @peer_threads[peer_spec][:peer] = resurrected_peer
          end
        end
      end
      Log.info("all peers exited, goodbye")
    end
  end
end

# Command line interface
#
if __FILE__ == $0
  include BGP

  options = {
    :logger => Logger.new(STDOUT),
    :local_pref => 100,
    :hold_time => 60,
    :bgp_identifier => nil,
    :autonomous_system => 65534,
    :routes => [],
    :peers =>  [],
    :communities => []
  }
  options[:logger].level = Logger::INFO

  OptionParser.new do |opts|
    opts.banner = "Usage: #{$0} [options...]"
    opts.separator ""
    opts.separator "Options (must specify as-number and bgp-identifier at least):"

    opts.on("-h", "--help", "Show this message") { print opts; exit }

    opts.on("-a", "--as-number NUM", Integer, "Set local AS number (default 65534)") do |as_num|
      options[:autonomous_system] = as_num
    end

    opts.on("-i", "--bgp-identifier ID", "Set BGP identifier") do |bgp_identifier|
      options[:bgp_identifier] = bgp_identifier
    end

    opts.on("-r", "--routes FILE", "Add routes file (default #{options[:routes]})") do |routes|
      options[:routes] << routes
    end

    opts.on("-p", "--peers FILE", "Add peers file (default #{options[:peers]})") do |peers|
      options[:peers] << peers
    end

    opts.on("--local-pref PREF", Integer, "Set local preference (default #{options[:local_pref]})") do |local_pref|
      options[:local_pref] = local_pref
    end

    opts.on("--hold-time SECS", Integer, "Set hold time (default #{options[:hold_time]})") do |hold_time|
      options[:hold_time] = hold_time
    end

    opts.on("--communities a,b,c", Array, "Set list of communities (default '#{options[:communities].join(",")}')") do |communities|
      options[:communities] = Update.parse_communities(communities)
    end

    opts.on("--log-level LEVEL", "Set log level (default INFO)") do |level|
      options[:logger].level = Logger.const_get(level)
    end

    opts.on("--log-file FILE", "Set log file (default STDOUT, or use SYSLOG)") do |file|
      options[:logger].close
      if file == 'SYSLOG'
        require 'syslog'
        Syslog.module_eval { class << self; alias :error :err; end } # hack!
        Syslog.module_eval { class << self; alias :warn :warning; end } # hack!
        Syslog.open('bgpfeeder')
        options[:logger] = Syslog
      else
        options[:logger] = Logger.new(file)
      end
    end

    opts.parse!(ARGV)
  end

  if !options[:bgp_identifier]
    STDERR.print "Must specify --bgp-identifier X (where X is your IP or a unique number in your AS)\n"
    exit 1
  end

  if options[:routes].empty? && options[:peers].empty?
    options[:routes] = ['routes']
    options[:peers] = ['peers']
  end

  if options[:routes].length != options[:peers].length
    STDERR.print "You must specify an equal number of --routes and --peers options\n"
    exit 1
  end

  BGP.const_set(:Log, options[:logger])
  options.delete(:logger)
  Log.info("Starting from command line with options: "+options.inspect)

  dbs = options[:routes].map do |route_file|
    db = Database.new(
                      options[:autonomous_system],
                      options[:bgp_identifier],
                      route_file
                      )
    db.communities = options[:communities]
    db.local_pref = options[:local_pref] if options[:local_pref]
    db.hold_time = options[:hold_time]
    db
  end

  peer_managers = []
  dbs.each_with_index do |db, i|
    peer_managers << PeerManager.new(options[:peers][i], db)
  end

  threads = []

  # first interrupt should kill all connections,
  trap("INT") do
    trap("INT","DEFAULT")
    peer_managers.each { |pm| pm.close_all }
  end

  threads = peer_managers.map { |pm| Thread.new { pm.run } }
  threads.each { |t| t.join }
end

