local ipOps = require "ipOps"
local math = require "math"
local nmap = require "nmap"
local packet = require "packet"
local stdnse = require "stdnse"
local string = require "string"
local table = require "table"

description = [[
Performs simple Path MTU Discovery to target hosts.

TCP or UDP packets are sent to the host with the DF (don't fragment) bit set
and with varying amounts of data.  If an ICMP Fragmentation Needed is received,
or no reply is received after retransmissions, the amount of data is lowered
and another packet is sent.  This continues until (assuming no errors occur) a
reply from the final host is received, indicating the packet reached the host
without being fragmented.

Not all MTUs are attempted so as to not expend too much time or network
resources.  Currently the relatively short list of MTUs to try contains
the plateau values from Table 7-1 in RFC 1191, "Path MTU Discovery".
Using these values significantly cuts down the MTU search space.  On top
of that, this list is rarely traversed in whole because:
* the MTU of the outgoing interface is used as a starting point, and
* we can jump down the list when an intermediate router sending a "can't fragment" message includes its next hop MTU (as described in RFC 1191 and required by RFC 1812)
]]

---
-- @usage
-- nmap --script path-mtu target
--
-- @output
-- Host script results:
-- |_path-mtu: 1492 <= PMTU < 1500
--
-- Host script results:
-- |_path-mtu: PMTU == 1006

author = "Kris Katterjohn"

license = "Same as Nmap--See https://nmap.org/book/man-legal.html"

categories = {"safe", "discovery"}


local IPPROTO_ICMP = packet.IPPROTO_ICMP
local IPPROTO_TCP  = packet.IPPROTO_TCP
local IPPROTO_UDP  = packet.IPPROTO_UDP

-- Number of times to retransmit for no reply before dropping to
-- another MTU value
local RETRIES = 1

-- RFC 1191, Table 7-1: Plateaus. Even the massive MTU values are
-- here since we skip down the list based on the outgoing interface
-- so its no harm.
local MTUS = {
  65535,
  32000,
  17914,
  8166,
  4352,
  2002,
  1492,
  1006,
  508,
  296,
  68
}

-- Find the index in MTUS{} to use based on the MTU +new+. If +new+ is in
-- between values in MTUS, then insert it into the table appropriately.
local searchmtu = function(cidx, new)
  if new == 0 then
    return cidx
  end

  while cidx <= #MTUS do
    if new >= MTUS[cidx] then
      if new ~= MTUS[cidx] then
        table.insert(MTUS, cidx, new)
      end
      return cidx
    end
    cidx = cidx + 1
  end
  return cidx
end

local dport = function(ip)
  if ip.ip_p == IPPROTO_TCP then
    return ip.tcp_dport
  elseif ip.ip_p == IPPROTO_UDP then
    return ip.udp_dport
  end
end

local sport = function(ip)
  if ip.ip_p == IPPROTO_TCP then
    return ip.tcp_sport
  elseif ip.ip_p == IPPROTO_UDP then
    return ip.udp_sport
  end
end

-- Checks how we should react to this packet
local checkpkt = function(reply, orig)
  local ip = packet.Packet:new(reply, reply:len())
  if not ip then
    return "recap"
  end

  if ip.ip_p == IPPROTO_ICMP then
    if ip.icmp_type ~= 3 then
      return "recap"
    end
    -- Port Unreachable
    if ip.icmp_code == 3 then
      local is = ip.buf:sub(ip.icmp_offset + 9)
      local ip2 = packet.Packet:new(is, is:len())

      -- Check sent packet against ICMP payload
      if not ip2 or ip2.ip_p ~= IPPROTO_UDP or
          ip2.ip_p ~= orig.ip_p or
          ip2.ip_bin_src ~= orig.ip_bin_src or
          ip2.ip_bin_dst ~= orig.ip_bin_dst or
          sport(ip2) ~= sport(orig) or
          dport(ip2) ~= dport(orig) then
        return "recap"
      end

      return "gotreply"
    end
    -- Frag needed, DF set
    if ip.icmp_code == 4 then
      local val = ip:u16(ip.icmp_offset + 6)
      return "nextmtu", val
    end
    return "recap"
  end

  if ip.ip_p ~= orig.ip_p or
      ip.ip_bin_src ~= orig.ip_bin_dst or
      ip.ip_bin_dst ~= orig.ip_bin_src or
      dport(ip) ~= sport(orig) or
      sport(ip) ~= dport(orig) then
    return "recap"
  end

  return "gotreply"
end

-- This is all we can use since we can get various protocols back from
-- different hosts
local check = function(layer3)
  local ip = packet.Packet:new(layer3, layer3:len())
  return ip and ip.ip_bin_dst
end

-- Updates a packet's info and calculates checksum
local updatepkt = function(ip)
  if ip.ip_p == IPPROTO_TCP then
    ip:tcp_set_sport(math.random(0x401, 0xffff))
    ip:tcp_set_seq(math.random(1, 0x7fffffff))
    ip:tcp_count_checksum()
  elseif ip.ip_p == IPPROTO_UDP then
    ip:udp_set_sport(math.random(0x401, 0xffff))
    ip:udp_set_length(ip.ip_len - ip.ip_hl * 4)
    ip:udp_count_checksum()
  end
  ip:ip_count_checksum()
end

-- Set up packet header and data to satisfy a certain MTU
local setmtu = function(pkt, mtu)
  if pkt.ip_len < mtu then
    pkt.buf = pkt.buf .. string.rep("\0", mtu - pkt.ip_len)
  else
    pkt.buf = pkt.buf:sub(1, mtu)
  end

  pkt:ip_set_len(mtu)
  pkt.packet_length = mtu
  updatepkt(pkt)
end

local basepkt = function(proto)
  local ibin = stdnse.fromhex(
    "4500 0014 0000 4000 8000 0000 0000 0000 0000 0000"
  )
  local tbin = stdnse.fromhex(
    "0000 0000 0000 0000 0000 0000 6002 0c00 0000 0000 0204 05b4"
  )
  local ubin = stdnse.fromhex(
    "0000 0000 0800 0000"
  )

  if proto == IPPROTO_TCP then
    return ibin .. tbin
  elseif proto == IPPROTO_UDP then
    return ibin .. ubin
  end
end

-- Creates a Packet object for the given proto and port
local genericpkt = function(host, proto, port)
  local pkt = basepkt(proto)
  local ip = packet.Packet:new(pkt, pkt:len())

  ip:ip_set_bin_src(host.bin_ip_src)
  ip:ip_set_bin_dst(host.bin_ip)

  ip:set_u8(ip.ip_offset + 9, proto)
  ip.ip_p = proto

  ip:ip_set_len(pkt:len())

  if proto == IPPROTO_TCP then
    ip:tcp_parse(false)
    ip:tcp_set_dport(port)
  elseif proto == IPPROTO_UDP then
    ip:udp_parse(false)
    ip:udp_set_dport(port)
  end

  updatepkt(ip)

  return ip
end

local ipproto = function(p)
  if p == "tcp" then
    return IPPROTO_TCP
  elseif p == "udp" then
    return IPPROTO_UDP
  end
  return -1
end

-- Determines how to probe
local getprobe = function(host)
  local combos = {
    { "tcp", "open" },
    { "tcp", "closed" },
    -- udp/open probably only happens when Nmap sends proper
    -- payloads, which doesn't happen in here
    { "udp", "closed" }
  }
  local proto = nil
  local port = nil

  for _, c in ipairs(combos) do
    port = nmap.get_ports(host, nil, c[1], c[2])
    if port then
      proto = c[1]
      break
    end
  end

  return proto, port
end

-- Sets necessary probe data in registry
local setreg = function(host, proto, port)
  host.registry['pathmtuprobe'] = {
    ['proto'] = proto,
    ['port'] = port
  }
end

hostrule = function(host)
  if not nmap.is_privileged() then
    nmap.registry[SCRIPT_NAME] = nmap.registry[SCRIPT_NAME] or {}
    if not nmap.registry[SCRIPT_NAME].rootfail then
      stdnse.verbose1("not running for lack of privileges.")
    end
    nmap.registry[SCRIPT_NAME].rootfail = true
    return nil
  end

  if nmap.address_family() ~= 'inet' then
    stdnse.debug1("is IPv4 compatible only.")
    return false
  end
  if not (host.interface and host.interface_mtu) then
    return false
  end
  local proto, port = getprobe(host)
  if not (proto and port) then
    return false
  end
  setreg(host, proto, port.number)
  return true
end

action = function(host)
  local m, r
  local gotit = false
  local mtuset
  local sock = nmap.new_dnet()
  local pcap = nmap.new_socket()
  local proto = host.registry['pathmtuprobe']['proto']
  local port = host.registry['pathmtuprobe']['port']
  local saddr = ipOps.str_to_ip(host.bin_ip_src)
  local daddr = ipOps.str_to_ip(host.bin_ip)
  local try = nmap.new_try()
  local status, pkt, ip

  try(sock:ip_open())

  try = nmap.new_try(function() sock:ip_close() end)

  pcap:pcap_open(host.interface, 104, false, "dst host " .. saddr .. " and (icmp or (" .. proto .. " and src host " .. daddr .. " and src port " .. port .. "))")

  -- Since we're sending potentially large amounts of data per packet,
  -- simply bump up the host's calculated timeout value.  Most replies
  -- should come from routers along the path, fragmentation reassembly
  -- times isn't an issue and the large amount of data is only traveling
  -- in one direction; still, we want a response from the target so call
  -- it 1.5*timeout to play it safer.
  pcap:set_timeout(1.5 * host.times.timeout * 1000)

  m = searchmtu(1, host.interface_mtu)

  mtuset = MTUS[m]

  local pkt = genericpkt(host, ipproto(proto), port)

  while m <= #MTUS do
    setmtu(pkt, MTUS[m])

    r = 0
    status = false
    while true do
      if not status then
        if not sock:ip_send(pkt.buf, host) then
          -- Got a send error, perhaps EMSGSIZE
          -- when we don't know our interface's
          -- MTU.  Drop an MTU and keep trying.
          break
        end
      end

      local test = pkt.ip_bin_src
      local status, length, _, layer3 = pcap:pcap_receive()
      while status and test ~= check(layer3) do
        status, length, _, layer3 = pcap:pcap_receive()
      end

      if status then
        local t, v = checkpkt(layer3, pkt)
        if t == "gotreply" then
          gotit = true
          break
        elseif t == "recap" then
        elseif t == "nextmtu" then
          if v == 0 then
            -- Router didn't send its
            -- next-hop MTU. Just drop
            -- a level.
            break
          end
          -- Lua's lack of a continue statement
          -- for loop control sucks, so dec m
          -- here as it's inc'd below.  Ugh.
          m = searchmtu(m, v) - 1
          mtuset = v
          break
        end
      else
        if r >= RETRIES then
          break
        end
        r = r + 1
      end
    end

    if gotit then
      break
    end

    m = m + 1
  end

  pcap:close()
  sock:ip_close()

  if not gotit then
    if nmap.debugging() > 0 then
      return "Error: Unable to determine PMTU (no replies)"
    end
    return
  end

  if MTUS[m] == mtuset then
    return "PMTU == " .. MTUS[m]
  elseif m == 1 then
    return "PMTU >= " .. MTUS[m]
  else
    return "" .. MTUS[m] .. " <= PMTU < " .. MTUS[m - 1]
  end
end
