local math = require "math"
local stdnse = require "stdnse"
local string = require "string"
local stringaux = require "stringaux"
local table = require "table"
local type     = type
local ipairs   = ipairs
local tonumber = tonumber
local unittest = require "unittest"
_ENV = stdnse.module("ipOps", stdnse.seeall)
local pack = string.pack
local unpack = string.unpack
isPrivate = function( ip )
  local err
  ip, err = expand_ip( ip )
  if err then return nil, err end
  if ip:match( ":" ) then
    local is_private
    local ipv6_private = { "::/127", "FC00::/7", "FE80::/10", "FEC0::/10" }
    for _, range in ipairs( ipv6_private ) do
      is_private, err = ip_in_range( ip, range )
      if is_private == true then
        return true, range
      end
      if err then
        return nil, err
      end
    end
  elseif ip:sub(1,3) == '10.' then
    return true, '10/8'
  elseif ip:sub(1,4) == '127.' then
    return true, '127/8'
  elseif ip:sub(1,8) == '169.254.' then
    return true, '169.254/16'
  elseif ip:sub(1,4) == '172.' then
    local p, e = ip_in_range(ip, '172.16/12')
    if p == true then
      return true, '172.16/12'
    else
      return p, e
    end
  elseif ip:sub(1,4) == '192.' then
    if ip:sub(5,8) == '168.' then
      return true, '192.168/16'
    elseif ip:match('^192%.[0][0]?[0]?%.[0][0]?[0]?%.') then
      return true, '192.0.0/24'
    elseif ip:match('^192%.[0][0]?[0]?%.[0]?[0]?2') then
      return true, '192.0.2/24'
    end
  elseif ip:sub(1,4) == '198.' then
    if ip:match('^198%.[0]?18%.') or ip:match('^198%.[0]?19%.') then
      return true, '198.18/15'
    elseif ip:match('^198%.[0]?51%.100%.') then
      return true, '198.51.100/24'
    end
  elseif ip:match('^203%.[0][0]?[0]?%.113%.') then
    return true, '203.0.113/24'
  elseif ip:match('^224%.[0][0]?[0]?%.[0][0]?[0]?%.') then
    return true, '224.0.0/24'
  elseif ip:match('^24[0-9]%.') or ip:match('^25[0-5]%.') then
    return true, '240.0.0/4'
  end
  return false, nil
end
todword = function( ip )
  if type( ip ) ~= "string" or ip:match( ":" ) then
    return nil, "Error in ipOps.todword: Expected IPv4 address."
  end
  local n, ret, err = {}
  n, err = get_parts_as_number( ip )
  if err then return nil, err end
  ret = (((n[1]*256+n[2]))*256+n[3])*256+n[4]
  return ret
end
fromdword = function( ip )
  if type( ip ) ~= "number" then
    stdnse.debug1("Error in ipOps.fromdword: Expected 32-bit number.")
    return nil
  end
  return string.format("%d.%d.%d.%d", unpack("BBBB", pack(">I4", ip)))
end
get_parts_as_number = function( ip )
  local err
  ip, err = expand_ip( ip )
  if err then return nil, err end
  local pattern, base
  if ip:match( ":" ) then
    pattern = "%x+"
    base = 16
  else
    pattern = "%d+"
    base = 10
  end
  local t = {}
  for part in string.gmatch(ip, pattern) do
    t[#t+1] = tonumber( part, base )
  end
  return t
end
compare_ip = function( left, op, right )
  if type( left ) ~= "string" or type( right ) ~= "string" then
    return nil, "Error in ipOps.compare_ip: Expected IP address as a string."
  end
  local err ={}
  left, err[#err+1] = ip_to_str( left )
  right, err[#err+1] = ip_to_str( right )
  if #err > 0 then
    return nil, table.concat( err, " " )
  end
  left = pack("s1", left)
  right = pack("s1", right)
  if ( op == "eq" ) then
    return ( left == right )
  elseif ( op == "ne" ) then
    return ( left ~= right )
  elseif ( op == "le" ) then
    return ( left <= right )
  elseif ( op == "ge" ) then
    return ( left >= right )
  elseif ( op == "lt" ) then
    return ( left < right )
  elseif ( op == "gt" ) then
    return ( left > right )
  end
  return nil, "Error in ipOps.compare_ip: Invalid Operator."
end
ip_sort = function (ips, op)
  op = op or "lt"
  return table.sort(ips, function(a, b) return compare_ip(a, op, b) end)
end
ip_in_range = function( ip, range )
  local first, last, err = get_ips_from_range( range )
  if err then return nil, err end
  ip, err = expand_ip( ip )
  if err then return nil, err end
  if ( ip:match( ":" ) and not first:match( ":" ) ) or ( not ip:match( ":" ) and first:match( ":" ) ) then
    return nil, "Error in ipOps.ip_in_range: IP address is of a different address family to Range."
  end
  err = {}
  local ip_ge_first, ip_le_last
  ip_ge_first, err[#err+1] = compare_ip( ip, "ge", first )
  ip_le_last, err[#err+1] = compare_ip( ip, "le", last )
  if #err > 0 then
    return nil, table.concat( err, " " )
  end
  if ip_ge_first and ip_le_last then
    return true
  else
    return false
  end
end
expand_ip = function( ip, family )
  local err
  if type( ip ) ~= "string" or ip == "" then
    return nil, "Error in ipOps.expand_ip: Expected IP address as a string."
  end
  local err4 = "Error in ipOps.expand_ip: An address assumed to be IPv4 was malformed."
  if not ip:match( ":" ) then
    if ip:match( "[^%.0-9]" ) then
      return nil, err4
    end
    local octets = {}
    for octet in string.gmatch( ip, "%d+" ) do
      if tonumber( octet, 10 ) > 255 then return nil, err4 end
      octets[#octets+1] = octet
    end
    if #octets > 4 then return nil, err4 end
    while #octets < 4 do
      octets[#octets+1] = "0"
    end
    if family == "inet6" then
      return ( table.concat( { 0,0,0,0,0,"ffff",
        stdnse.tohex( 256*octets[1]+octets[2] ),
        stdnse.tohex( 256*octets[3]+octets[4] )
        }, ":" ) )
    else
      return ( table.concat( octets, "." ) )
    end
  end
  if family ~= nil and family ~= "inet6" then
    return nil, "Error in ipOps.expand_ip: Cannot convert IPv6 address to IPv4"
  end
  if ip:match( "[^%.:%x]" ) then
    return nil, ( err4:gsub( "IPv4", "IPv6" ) )
  end
  ip = string.gsub(ip, "::", ":z:")
  local hexadectets = {}
  for hdt in string.gmatch( ip, "[%.z%x]+" ) do
    hexadectets[#hexadectets+1] = hdt
  end
  local t = {}
  if hexadectets[#hexadectets]:match( "[%.]+" ) then
    hexadectets[#hexadectets], err = expand_ip( hexadectets[#hexadectets] )
    if err then return nil, ( err:gsub( "IPv4", "IPv4in6" ) ) end
    t = stringaux.strsplit( "[%.]+", hexadectets[#hexadectets] )
    for i, v in ipairs( t ) do
      t[i] = tonumber( v, 10 )
    end
    hexadectets[#hexadectets] = stdnse.tohex( 256*t[1]+t[2] )
    hexadectets[#hexadectets+1] = stdnse.tohex( 256*t[3]+t[4] )
  end
  local z_done = false
  for index, value in ipairs( hexadectets ) do
    if value:match( "[%.]+" ) then
      return nil, ( err4:gsub( "IPv4", "IPv6" ) )
    elseif value == "z" and z_done then
      return nil, ( err4:gsub( "IPv4", "IPv6" ) )
    elseif value == "z" and not z_done then
      z_done = true
      hexadectets[index] = "0"
      local bound = 8 - #hexadectets
      for i = 1, bound, 1 do
        table.insert( hexadectets, index+i, "0" )
      end
    elseif tonumber( value, 16 ) > 65535 then
      return nil, ( err4:gsub( "IPv4", "IPv6" ) )
    end
  end
  if #hexadectets > 8 then return nil, ( err4:gsub( "IPv4", "IPv6" ) ) end
  while #hexadectets < 8 do
    hexadectets[#hexadectets+1] = "0"
  end
  return ( table.concat( hexadectets, ":" ) )
end
get_ips_from_range = function( range )
  if type( range ) ~= "string" then
    return nil, nil, "Error in ipOps.get_ips_from_range: Expected a range as a string."
  end
  local ip, prefix = range:match("^%s*([%x:.]+)/(%d+)%s*$")
  if ip then
    return get_first_last_ip(ip, prefix)
  end
  local first, last = range:match("^%s*([%x:.]+)%s*%-%s*([%x:.]+)%s*$")
  if not first then
    return nil, nil, "Error in ipOps.get_ips_from_range: The range supplied could not be interpreted."
  end
  local err
  first, err = expand_ip(first)
  if not err then
    last, err = expand_ip(last)
  end
  if not err then
    local af = function (ip) return ip:find(":") and 6 or 4 end
    if af(first) ~= af(last) then
      err = "Error in ipOps.get_ips_from_range: First IP address is of a different address family to last IP address."
    end
  end
  if err then
      return nil, nil, err
  end
  return first, last
end
get_first_last_ip = function(ip, prefix)
  local err
  ip, err = ip_to_bin(ip)
  if err then return nil, nil, err end
  prefix = tonumber(prefix)
  if not prefix or prefix ~= math.floor(prefix) or prefix < 0 or prefix > #ip then
    return nil, nil, "Error in ipOps.get_first_last_ip: Invalid prefix."
  end
  local net = ip:sub(1, prefix)
  local first = bin_to_ip(net .. ("0"):rep(#ip - prefix))
  local last  = bin_to_ip(net .. ("1"):rep(#ip - prefix))
  return first, last
end
get_first_ip = function(ip, prefix)
  local first, last, err = get_first_last_ip(ip, prefix)
  return first, err
end
get_last_ip = function(ip, prefix)
  local first, last, err = get_first_last_ip(ip, prefix)
  return last, err
end
ip_to_str = function( ip, family )
  local err
  ip, err = expand_ip( ip, family )
  if err then return nil, err end
  local t = {}
  if not ip:match( ":" ) then
    for octet in string.gmatch( ip, "%d+" ) do
      t[#t+1] = pack("B", octet)
    end
  else
    for hdt in string.gmatch( ip, "%x+" ) do
      t[#t+1] = pack( ">I2", tonumber(hdt, 16) )
    end
  end
  return table.concat( t )
end
str_to_ip = function (ip)
  if #ip == 4 then
    return ("%d.%d.%d.%d"):format(unpack("BBBB", ip))
  elseif #ip == 16 then
    local full = ("%x:%x:%x:%x:%x:%x:%x:%x"):format(unpack((">I2"):rep(8), ip))
    full = full:gsub(":[:0]+", "::", 1) -- Collapse the first (should be longest?) series of :0:
    full = full:gsub("^0::", "::", 1) -- handle special case of ::1
    return full
  else
    return nil, "Invalid length"
  end
end
ip_to_bin = function( ip )
  local err
  ip, err = expand_ip( ip )
  if err then return nil, err end
  local t, mask = {}
  if not ip:match( ":" ) then
    for octet in string.gmatch( ip, "%d+" ) do
      t[#t+1] = stdnse.tohex( tonumber(octet) )
    end
    mask = "00"
  else
    for hdt in string.gmatch( ip, "%x+" ) do
      t[#t+1] = hdt
    end
    mask = "0000"
  end
  for i, v in ipairs( t ) do
    t[i] = mask:sub( 1, # mask  - # v  ) .. v
  end
  return hex_to_bin( table.concat( t ) )
end
bin_to_ip = function( binstring )
  if type( binstring ) ~= "string" or binstring:match( "[^01]+" ) then
    return nil, "Error in ipOps.bin_to_ip: Expected string of binary digits."
  end
  local af
  if # binstring  == 32 then
    af = 4
  elseif # binstring  == 128 then
    af = 6
  else
    return nil, "Error in ipOps.bin_to_ip: Expected exactly 32 or 128 binary digits."
  end
  local t = {}
  if af == 6 then
    local pattern = string.rep( "[01]", 16 )
    for chunk in string.gmatch( binstring, pattern ) do
      t[#t+1] = stdnse.tohex( tonumber( chunk, 2 ) )
    end
    return table.concat( t, ":" )
  end
  if af == 4 then
    local pattern = string.rep( "[01]", 8 )
    for chunk in string.gmatch( binstring, pattern ) do
      t[#t+1] = tonumber( chunk, 2 ) .. ""
    end
    return table.concat( t, "." )
  end
end
local bin_lookup = {
  [0]="0000", "0001", "0010", "0011", "0100", "0101", "0110", "0111",
      "1000", "1001", "1010", "1011", "1100", "1101", "1110", "1111",
}
hex_to_bin = function( hex )
  if type( hex ) ~= "string" then
    return nil, "Error in ipOps.hex_to_bin: Expected string"
  end
  local status, result = pcall( string.gsub, hex, ".", function(nibble)
      local n = bin_lookup[tonumber(nibble, 16)]
      if n then
        return n
      else
        error("Error in ipOps.hex_to_bin: Expected string representing a hexadecimal number.")
      end
    end)
  if status then
    return result
  end
  return status, result
end
cidr_to_subnet = function( subnet )
  local bits = subnet:match("/(%d%d)$")
  if not bits then return nil end
  return fromdword((0xFFFFFFFF >> tonumber(bits)) ~ 0xFFFFFFFF)
end
subnet_to_cidr = function( subnet )
  local dword, err = todword(subnet)
  if not dword then return nil, err end
  return "/" .. tostring(32 - (math.tointeger(math.log((dword ~ 0xFFFFFFFF) + 1, 2))))
end
if not unittest.testing() then
  return _ENV
end
test_suite = unittest.TestSuite:new()
test_suite:add_test(unittest.is_true(isPrivate("192.168.123.123")), "192.168.123.123 is private")
test_suite:add_test(unittest.is_false(isPrivate("1.1.1.1")), "1.1.1.1 is not private")
test_suite:add_test(unittest.equal(todword("65.66.67.68"),0x41424344), "todword")
test_suite:add_test(unittest.equal(todword("127.0.0.1"),0x7f000001), "todword")
test_suite:add_test(unittest.equal(fromdword(0xffffffff),"255.255.255.255"), "fromdword")
test_suite:add_test(unittest.equal(fromdword(0x7f000001),"127.0.0.1"), "fromdword")
test_suite:add_test(unittest.equal(str_to_ip("\x01\x02\x03\x04"),"1.2.3.4"), "str_to_ip (ipv4)")
test_suite:add_test(unittest.equal(str_to_ip("\0\x01\xbe\xef\0\0\0\0\0\0\x02\x03\0\0\0\x01"),"1:beef::203:0:1"), "str_to_ip (ipv6)")
test_suite:add_test(unittest.equal(str_to_ip(("\0"):rep(15) .. "\x01"),"::1"), "str_to_ip (ipv6)")
test_suite:add_test(function()
  local parts, err = get_parts_as_number("8.255.0.1")
  if parts == nil then return false, err end
  if parts[1] == 8 and parts[2] == 255 and parts[3] == 0 and parts[4] == 1 then
    return true
  end
  return false, string.format("Expected {8, 255, 0, 1}, got {%d, %d, %d, %d}", table.unpack(parts))
end, "get_parts_as_number")
do
  local low_ip4 = "192.168.1.10"
  local high_ip4 = "192.168.10.1"
  local low_ip6 = "2001::DEAD:0:0:9"
  local high_ip6 = "2001::DEAF:0:0:9"
  for _, op in ipairs({
    {low_ip4, "eq", low_ip4, unittest.is_true, "IPv4"},
    {low_ip6, "eq", low_ip6, unittest.is_true, "IPv6"},
    {high_ip4, "eq", low_ip4, unittest.is_false, "IPv4"},
    {high_ip6, "eq", low_ip6, unittest.is_false, "IPv6"},
    {low_ip4, "eq", low_ip6, unittest.is_false, "mixed"},
    {low_ip4, "ne", low_ip4, unittest.is_false, "IPv4"},
    {low_ip6, "ne", low_ip6, unittest.is_false, "IPv6"},
    {high_ip4, "ne", low_ip4, unittest.is_true, "IPv4"},
    {high_ip6, "ne", low_ip6, unittest.is_true, "IPv6"},
    {low_ip4, "ne", low_ip6, unittest.is_true, "mixed"},
    {low_ip4, "ge", low_ip4, unittest.is_true, "IPv4, equal"},
    {low_ip6, "ge", low_ip6, unittest.is_true, "IPv6, equal"},
    {high_ip4, "ge", low_ip4, unittest.is_true, "IPv4"},
    {high_ip6, "ge", low_ip6, unittest.is_true, "IPv6"},
    {low_ip4, "ge", high_ip4, unittest.is_false, "IPv4"},
    {low_ip6, "ge", high_ip6, unittest.is_false, "IPv6"},
    {low_ip6, "ge", low_ip4, unittest.is_true, "mixed"},
    {low_ip4, "ge", low_ip6, unittest.is_false, "mixed"},
    {low_ip4, "le", low_ip4, unittest.is_true, "IPv4, equal"},
    {low_ip6, "le", low_ip6, unittest.is_true, "IPv6, equal"},
    {high_ip4, "le", low_ip4, unittest.is_false, "IPv4"},
    {high_ip6, "le", low_ip6, unittest.is_false, "IPv6"},
    {low_ip4, "le", high_ip4, unittest.is_true, "IPv4"},
    {low_ip6, "le", high_ip6, unittest.is_true, "IPv6"},
    {low_ip6, "le", low_ip4, unittest.is_false, "mixed"},
    {low_ip4, "le", low_ip6, unittest.is_true, "mixed"},
    {low_ip4, "gt", low_ip4, unittest.is_false, "IPv4, equal"},
    {low_ip6, "gt", low_ip6, unittest.is_false, "IPv6, equal"},
    {high_ip4, "gt", low_ip4, unittest.is_true, "IPv4"},
    {high_ip6, "gt", low_ip6, unittest.is_true, "IPv6"},
    {low_ip4, "gt", high_ip4, unittest.is_false, "IPv4"},
    {low_ip6, "gt", high_ip6, unittest.is_false, "IPv6"},
    {low_ip6, "gt", low_ip4, unittest.is_true, "mixed"},
    {low_ip4, "gt", low_ip6, unittest.is_false, "mixed"},
    {low_ip4, "lt", low_ip4, unittest.is_false, "IPv4, equal"},
    {low_ip6, "lt", low_ip6, unittest.is_false, "IPv6, equal"},
    {high_ip4, "lt", low_ip4, unittest.is_false, "IPv4"},
    {high_ip6, "lt", low_ip6, unittest.is_false, "IPv6"},
    {low_ip4, "lt", high_ip4, unittest.is_true, "IPv4"},
    {low_ip6, "lt", high_ip6, unittest.is_true, "IPv6"},
    {low_ip6, "lt", low_ip4, unittest.is_false, "mixed"},
    {low_ip4, "lt", low_ip6, unittest.is_true, "mixed"},
    }) do
    test_suite:add_test(op[4](compare_ip(op[1], op[2], op[3])),
      string.format("compare_ip(%s, %s, %s) (%s)", op[1], op[2], op[3], op[5]))
  end
end
do
  for _, h in ipairs({
      {"a", "1010"},
      {"aa", "10101010"},
      {"12", "00010010"},
      {"54321", "01010100001100100001"},
      {"123error", false},
      {"", ""},
      {"bad 123", false},
    }) do
    test_suite:add_test(unittest.equal(hex_to_bin(h[1]), h[2]))
  end
end
do
  for _, op in ipairs({
    {"192.168.13.1", "192/8", unittest.is_true, "IPv4 CIDR"},
    {"193.168.13.1", "192/8", unittest.is_false, "IPv4 CIDR"},
    {"192.168.13.0", "192.168.13.128/24", unittest.is_true, "IPv4 CIDR"},
    {"193.168.13.0", "192.168.13.128/24", unittest.is_false, "IPv4 CIDR"},
    {"2001:db8::9", "2001:db8/32", unittest.is_true, "IPv6 CIDR"},
    {"2001:db7::9", "2001:db8/32", unittest.is_false, "IPv6 CIDR"},
    {"2001:db8::9", "2001:db8::1:0/32", unittest.is_true, "IPv6 CIDR"},
    {"2001:db7::9", "2001:db8::1:0/32", unittest.is_false, "IPv6 CIDR"},
    {"192.168.13.1", "192.168.10.33-192.168.80.80", unittest.is_true, "IPv4 range"},
    {"193.168.13.1", "192.168.1.1 - 192.168.5.0", unittest.is_false, "IPv4 range"},
    {"2001:db8::9", "2001:db8::1-2001:db8:1::1", unittest.is_true, "IPv6 range"},
    {"2001:db8::9", "2001:db8:10::1-2001:db8:11::1", unittest.is_false, "IPv6 range"},
    {"193.168.1.1", "192.168.1.1 - 2001:db8::1", unittest.is_nil, "mixed"},
    {"2001:db8::1", "192.168.1.1 - 2001:db8::1", unittest.is_nil, "mixed"},
    }) do
    test_suite:add_test(op[3](ip_in_range(op[1], op[2])),
      string.format("ip_in_range(%s, %s) (%s)", op[1], op[2], op[4]))
  end
end
do
  for _, op in ipairs({
    {"192.168", nil, "192.168.0.0", "IPv4 trunc"},
    {"192.0.2.3", nil, "192.0.2.3", "IPv4"},
    {"192.168", "inet6", "0:0:0:0:0:ffff:c0a8:0", "IPv4 trunc to IPv6"},
    {"2001:db8::9", nil, "2001:db8:0:0:0:0:0:9", "IPv6"},
    {"::ffff:192.0.2.128", "inet6", "0:0:0:0:0:ffff:c000:280", "IPv4-mapped to IPv6"},
    }) do
    test_suite:add_test(unittest.equal(expand_ip(op[1], op[2]), op[3]),
      string.format("expand_ip(%s, %s) (%s)", op[1], op[2], op[4]))
  end
  test_suite:add_test(unittest.is_nil(expand_ip("2001:db8::1", "ipv4")),
      "IPv6 to IPv4")
end
test_suite:add_test(unittest.equal(cidr_to_subnet("/16"), "255.255.0.0"), "cidr_to_subnet")
test_suite:add_test(unittest.equal(subnet_to_cidr("255.255.0.0"), "/16"), "subnet_to_cidr")
return _ENV;
