#!/usr/bin/ruby
# NATPMP -- a Ruby module for NAT-PMP.
# See http://files.dns-sd.org/draft-cheshire-nat-pmp.txt
#
# Copyright (c) 2007 Evan Martin <martine@danga.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

require 'socket'

module NATPMP
  PMP_PORT = 5351

class ProtocolException < RuntimeError; end

class Request
  PMP_VERSION = 0

  def initialize(*args)
    @encoded = encode(*args)
  end

  def encode(op)
    [PMP_VERSION, op].pack('CC')
  end

  def response_size; 8; end

  def decode_from(data)
    @version, @op, @result, @uptime = data.unpack('CCnN')
    @op -= 128
    raise ProtocolException, result_as_str(@result) unless @result == 0
    data[8..-1]
  end

  def result_as_str(result)
    case result
    when 0; "Success"
    when 1; "Unsupported Version"
    when 2; "Not Authorized/Refused (e.g. box supports mapping, but user has turned feature off)"
    when 3; "Network Failure (e.g. NAT box itself has not obtained a DHCP lease)"
    when 4; "Out of resources (NAT box cannot create any more mappings at this time)"
    when 5; "Unsupported opcode"
    else;   "Unknown result code (NAT-PMP protocol violation?)"
    end
  end

  def dump
    <<-EOT
version: #{@version}
op:      #{@op}
result:  #{@result}: #{result_as_str(@result)}
uptime:  #{@uptime}
    EOT
end

  def run(sock)
    sock.send(@encoded, 0)
    data, sender = sock.recvfrom(response_size)
    decode_from(data)
    return self
  end
end

class IPRequest < Request
  def initialize
    super(0)
  end

  def response_size; super + 4; end

  def decode_from(data)
    data = super(data)
    @ip = data.unpack('N')[0]
    data[4..-1]
  end

  def dump
    super + <<-EOT
ip:      #{addr_str}
    EOT
  end

  def addr_str
    addr = [@ip].pack('N').unpack('C4')
    return "#{addr[0]}.#{addr[1]}.#{addr[2]}.#{addr[3]}"
  end
end

class MappingRequest < Request
  def encode(tcp, localport, publicport, seconds)
    super(tcp ? 2 : 1) + [0, localport, publicport, seconds].pack('nnnN')
  end
  def response_size; super + 2*4; end
  def decode_from(data)
    data = super(data)
    @privport, @pubport, @seconds = data.unpack('nnN')
    data[8..-1]
  end
  def dump
    super + <<-EOT
privport: #{@privport}
pubport:  #{@pubport}
seconds:  #{@pubport}
    EOT
  end
end

end  # module PMP

if __FILE__ == $0
  sock = UDPSocket.new
  sock.connect('10.0.0.1', NATPMP::PMP_PORT)
  ip = NATPMP::IPRequest.new.run(sock)
  puts "Local IP request:"
  puts ip.dump
  puts

  # Punch a hole in for ssh to my machine for ten minutes.
  puts "Port mapping request:"
  map = NATPMP::MappingRequest.new(true, 22, 23, 60*10)
  map.run(sock)
  puts map.dump
end