/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

// General functions for network address parsing and manipulation
package netaddr

import (
	"bytes"
	"fmt"
	"io"
	"math"
	"net"
	"strconv"
)

func byteOp(u, v []byte, f func(byte, byte) byte) []byte {
	w := make([]byte, len(u))

	for i := 0; i < len(u); i++ {
		w[i] = f(u[i], v[i])
	}
	return w
}

func byteOr(u, v []byte) []byte {
	return byteOp(u, v, func(ui, vi byte) byte {
		return ui | vi
	})
}

func byteXor(u, v []byte) []byte {
	return byteOp(u, v, func(ui, vi byte) byte {
		return ui ^ vi
	})
}

func applyMask(address, mask []byte) {
	for i := range address {
		address[i] &= mask[i]
	}
	return
}

var rfc1918 = []*NetAddr{
	&NetAddr{[]byte{10, 0, 0, 0}, uint8(8)},
	&NetAddr{[]byte{172, 16, 0, 0}, uint8(12)},
	&NetAddr{[]byte{192, 168, 0, 0}, uint8(16)},
}

func NewNetAddr(a []byte, l int) *NetAddr {
	addr := make([]byte, len(a))
	copy(addr, a)
	return &NetAddr{
		address: addr,
		length:  uint8(clamp(l, 0, len(a)*8)),
	}
}

/*
NetAddr stores an address of S bytes together with a prefix
length L.

A prefix length L < S induces a partition on the entire
address space, containing 2^L equivalence classes, each with
2^(S-L) elements. The canonical element for a given
equivalence class is always the least (or lowest) address.

However, as any element of the equivalence class is
representative of the rest, we allow the host bits to vary,
regardless of the prefix length. As this is not proper CIDR,
it is useful to think of this construction as an address
with context.
*/
type NetAddr struct {
	address []byte
	length  uint8
}

/* Address returns an address without context (length). */
func (na *NetAddr) Address() []byte {
	return na.address
}

/* Length returns the prefix length. */
func (na *NetAddr) Length() int {
	return int(na.length)
}

/* AddressSize returns the size of the address in bytes. */
func (na *NetAddr) AddressSize() int {
	return len(na.Address())
}

/*
SetLength changes the length.

The NetAddr pointer is returned for convenience.
*/
func (na *NetAddr) SetLength(length int) *NetAddr {
	na.length = uint8(min(length, na.AddressSize()*8))
	return na
}

/*
FirstAddress sets the address to the first (lowest)
address in the prefix.

The NetAddr pointer is returned for convenience.
*/
func (na *NetAddr) FirstAddress() *NetAddr {
	bitSize := na.AddressSize() * 8
	mask := net.CIDRMask(na.Length(), bitSize)
	applyMask(na.address, mask)
	return na
}

/*
LastAddress sets the address to the last (highest) address
in the prefix.

The NetAddr pointer is returned for convenience.
*/
func (na *NetAddr) LastAddress() *NetAddr {
	bitSize := na.AddressSize() * 8
	mask := net.CIDRMask(na.Length(), bitSize)
	allOnes := net.CIDRMask(bitSize, bitSize)
	invMask := byteXor(mask, allOnes)
	na.address = byteOr(na.FirstAddress().address, invMask)
	return na
}

/*
Clone returns a pointer to a new (deep) copy of the
provided NetAddr.
*/
func (na *NetAddr) Clone() *NetAddr {
	next := &NetAddr{
		address: make([]byte, len(na.address)),
		length:  na.length,
	}
	copy(next.address, na.address)
	return next
}

/*
Mask returns a dotted decimal subnet mask corresponding
to the given NetAddr prefix length.
*/
func (na *NetAddr) Mask() []byte {
	onesToMask :=
		[9]byte{0, 128, 192, 224, 240, 248, 252, 254, 255}
	ones := na.Length()
	mask := make([]byte, na.AddressSize())
	for i := range mask {
		olen := min(8, ones)
		ones -= olen
		mask[i] = onesToMask[int(olen)]
	}
	return mask
}

/*
Integer returns an aton-style integer representation for the
address of the given NetAddr.
*/
func (na *NetAddr) Integer() (dec uint64) {
	for _, b := range na.address {
		dec <<= 8
		dec += uint64(b)
	}
	return
}

/*
String returns an appopriate string representation for
the given NetAddr.
*/
func (na *NetAddr) String() string {
	addr := net.IP(na.address).String()
	length := fmt.Sprintf("%d", na.length)
	return addr + "/" + length
}

/*
IsRFC1918 returns true when the given NetAddr is a subset of
an RFC1918 reserved range.
*/
func (na *NetAddr) IsRFC1918() bool {
	for _, e := range rfc1918 {
		if na.IsSubset(e) {
			return true
		}
	}
	return false
}

/*
IsEqual returns true when the given NetAddrs have the same
address and length.
*/
func (na1 *NetAddr) IsEqual(na2 *NetAddr) bool {
	return NetAddrIsEqual(na1, na2)
}

/*
Contains returns true when na1 contains, or is equal to,
na2.
*/
func (na1 *NetAddr) Contains(na2 *NetAddr) bool {
	return NetAddrContains(na1, na2)
}

/*
IsSubset returns true when na2 contains, or is equal to,
na1.
*/
func (na1 *NetAddr) IsSubset(na2 *NetAddr) bool {
	return NetAddrIsSubset(na1, na2)
}

/*
Overlaps returns true when na1 Contains na2 or na2 Contains
na1.
*/
func (na1 *NetAddr) Overlaps(na2 *NetAddr) bool {
	return NetAddrOverlaps(na1, na2)
}

/* Test against sheared isometry on byte prefixes. */
func octetLessThan(o1, m1, o2, m2 byte) bool {
	/* We seek an isometry from a binary trie embedded in R²
	* to the integral number line, with ones bits branching to
	* the right and zeros bits branching to the left. When
	* considering only the ones bits of an n-bit binary
	* number, the usual decimal encoding suffices, as each ith
	* bit, for 0 <= i <= n-1, shifts a point at the origin to
	* the right by 2^i. However, the zeros bits have no
	* effect, in this sense, as leading zeros do not
	* contribute to the integer value, resulting in many
	* binary strings occupying the same point. This mapping is
	* well-formed and surjective (every integer has a binary
	* encoding), but it is not injective.
	*
	* -> R 0 00 000 0000
	*	    			 -> 0001
	*	    	 -> 001 0010
	*	    			 -> 0011
	*	    -> 01 010 0100
	*	    			 -> 0101
	*	    	 -> 011 0110
	*	    			 -> 0111
	*	  -> 1 10 100 1000
	*	    			 -> 1001
	*	    	 -> 101 1010
	*	    			 -> 1011
	*	    -> 11 110 1100
	*	    			 -> 1101
	*	    	 -> 111 1110
	*	    			 -> 1111
	*
	* Instead, by allowing each zero bit to provide a negative
	* contribution, we ensure that every path (binary string)
	* maps to a unique point on the integral line. To see that
	* this is an isometry, project the nodes of a binary trie
	* embedded in R², with branch width w = 2^(n-j-1) at rank
	* j and arbitrary branch height h, onto the integral line.
	*
	* Here, we express the above isometry by simply
	* subtracting the contributions of the place values of the
	* zeros bits, which we obtain by ones complement (XOR).
	* Masking ensures we map only the prefix.
	*
	* Finally, we note that projecting the usual geometry of a
	* binary trie onto the integral line places each ancestor
	* at the midpoint of its descendants. This creates a
	* disagreement between the partial and total orders, which
	* makes it difficult to recover the original partial order
	* generated by the trie.
	*
	*                   -> 0000
	*             -> 000
	*                   -> 0001
	*        -> 00
	*                   -> 0010
	*             -> 001
	*                   -> 0011
	*    -> 0
	*                   -> 0100
	*             -> 010
	*                   -> 0101
	*        -> 01
	*                   -> 0110
	*             -> 011
	*                   -> 0111
	* 0/0
	*                   -> 1000
	*             -> 100
	*                   -> 1001
	*        -> 10
	*                   -> 1010
	*             -> 101
	*                   -> 1011
	*    -> 1
	*                   -> 1100
	*             -> 110
	*                   -> 1101
	*        -> 11
	*                   -> 1110
	*             -> 111
	*                   -> 1111
	*
	* Instead, we require that the total order < given by the
	* proposed isometry have the property that it always agree
	* with the partial order ≪ given by the binary trie.  That
	* is, given nodes a, b of a binary trie, if a ≪ b, then a
	* < b. Note that the converse does not hold, nor would it
	* be useful if it did. To accomplish this, we effectively
	* apply a horizontal shear to the right on the binary trie
	* such that the root of each subtree falls to the right of
	* its descendants.
	*
	*                   -> 0000
	*                   -> 0001
	*             -> 000
	*                   -> 0010
	*                   -> 0011
	*             -> 001
	*        -> 00
	*                   -> 0100
	*                   -> 0101
	*             -> 010
	*                   -> 0110
	*                   -> 0111
	*             -> 011
	*        -> 01
	*    -> 0
	*                   -> 1000
	*                   -> 1001
	*             -> 100
	*                   -> 1010
	*                   -> 1011
	*             -> 101
	*        -> 10
	*                   -> 1100
	*                   -> 1101
	*             -> 110
	*                   -> 1110
	*                   -> 1111
	*             -> 111
	*        -> 11
	*    -> 1
	* 0/0
	*
	* In the unsheared isometry, the nodes we wish to reorder
	* are specifically those which precede their descendants--
	* that is, whenever a < b and b ≪ a, we would like to
	* move a to the right of b so that b < a, thereby agreeing
	* with the partial order. Since b ≪ a, it is the case
	* that a represents a shorter prefix than b. Furthermore,
	* since b is a descendant of a, b's prefix is guaranteed
	* to share all of the bits of a. Conversely, if a and b
	* share a prefix, then that prefix represents the most
	* recent shared ancestor (least upper bound). And if the
	* prefix shared by a and b is a's prefix, then it must be
	* the case that a's prefix is shorter, and b is a
	* descendant of a. Thus, b ≪ a.
	*
	* Therefore, we conclude that a must be reordered
	* precisely when a < b and b shares a's prefix. Hence our
	* assertion `f1 < f2 && o1 != o2&m1`. To reorder, we
	* simply place prefixes with longer masks before shorter
	* ones, as given by `m1 > m2`.
	**/
	f1 := int(o1&m1) - int((o1^byte(255))&m1)
	f2 := int(o2&m2) - int((o2^byte(255))&m2)
	return f1 < f2 && o1 != o2&m1 || m1 > m2
}

/*
LessThan imposes a total order < on NetAddrs, returning
true precisely when na1 < na2.
*/
func LessThan(na1 *NetAddr, na2 *NetAddr) (less bool) {
	if na1.AddressSize() != na2.AddressSize() {
		return
	}
	mask1 := na1.Mask()
	mask2 := na2.Mask()
	for i := 0; i < na1.AddressSize(); i++ {
		if octetLessThan(
			na1.address[i], mask1[i],
			na2.address[i], mask2[i],
		) {
			less = true
			break
		}
		if na1.address[i] == na2.address[i] &&
			mask1[i] == mask2[i] {
			continue
		}
		break
	}
	return
}

/*
FirstAddress is a non-mutating version of the FirstAddress
method.
*/
func FirstAddress(na *NetAddr) *NetAddr {
	return na.Clone().FirstAddress()
}

/*
LastAddress is a non-mutating version of the LastAddress
method.
*/
func LastAddress(na *NetAddr) *NetAddr {
	return na.Clone().LastAddress()
}

/*
NetAddrIsEqual is a function implementation of the IsEqual
method.
*/
func NetAddrIsEqual(na1, na2 *NetAddr) bool {
	return na1.length == na2.length &&
		bytes.Equal(na1.address, na2.address)
}

/*
NetAddrContains is a function implementation of the
Contains method.
*/
func NetAddrContains(na1, na2 *NetAddr) bool {
	tmp := na2.Clone().SetLength(na1.Length())
	return na1.length <= na2.length &&
		FirstAddress(na1).IsEqual(FirstAddress(tmp))
}

/*
NetAddrIsSubset is a function implementation of the IsSubset
method.
*/
func NetAddrIsSubset(na1, na2 *NetAddr) bool {
	return na2.Contains(na1)
}

/*
NetAddrOverlaps is a function implementation of the Overlaps
method.
*/
func NetAddrOverlaps(na1, na2 *NetAddr) bool {
	return na1.Contains(na2) || na2.Contains(na1)
}

/*
Contiguous tests whether na1 and na2 can be summarized by
a single prefix containing only na1 and na2.
*/
func Contiguous(na1, na2 *NetAddr) bool {
	if na1.AddressSize() != na2.AddressSize() ||
		na1.Length() != na2.Length() ||
		Join(na1, na2).Length() != na1.Length()-1 {
		return false
	}
	return true
}

/*
LeftAdjacent tests whether na1 is left-adjacent to na2.

LeftAdjacent returns true if, and only if, the first address
in the prefix of na2 is immediately preceded by the last
address in the prefix of na1. If na1 and na2 are contiguous,
then either na1 is adjacent to na2 or na2 is adjacent to
na1. However, neither converse holds, in general.
*/
func LeftAdjacent(na1, na2 *NetAddr) bool {
	if na1.AddressSize() != na2.AddressSize() {
		return false
	}
	b1 := na1.LastAddress().Integer()
	a2 := na2.FirstAddress().Integer()
	return b1+1 == a2
}

/*
	Join returns the smallest summary containing both na1 and

na2.
*/
func Join(na1, na2 *NetAddr) *NetAddr {
	minLen := min(na1.Length(), na2.Length())
	addrSize := min(na1.AddressSize(), na2.AddressSize())
	bitSize := addrSize * 8
	fa1 := na1.FirstAddress()
	fa2 := na2.FirstAddress()
	var bits int
	for i := 0; i < addrSize; i++ {
		xor := fa1.address[i] ^ fa2.address[i]
		if xor != 0 {
			bits += int(math.Trunc(1.0 + math.Log2(float64(xor))))
			bits += bitSize - 8*(i+1)
			break
		}
	}
	return fa1.
		Clone().
		SetLength(min(minLen, bitSize-bits)).
		FirstAddress()
}

/*
Split returns the two halves of a given NetAddr. The input
is unchanged.

If the given NetAddr represents a single address, then
Split returns the input for both the left and right outputs.
*/
func Split(na *NetAddr) (*NetAddr, *NetAddr) {
	if na.Length() == na.AddressSize()*8 {
		return na, na
	}
	nextLen := na.Length() + 1
	lower :=
		na.
			Clone().
			FirstAddress().
			SetLength(nextLen)
	upper :=
		na.
			Clone().
			LastAddress().
			SetLength(nextLen).
			FirstAddress()
	return lower, upper
}

func parseTruncatedIPv4(s string) (na *NetAddr, err error) {
	na = &NetAddr{address: make([]byte, 4)}
	r := bytes.NewBufferString(s)
	buf := new(bytes.Buffer)
	i := 0
	onLen := false
	for i < 4 {
		var b byte
		b, err = r.ReadByte()
		switch {
		case err == io.EOF:
			na.length = byte(8*i - 8)
			if onLen {
				var o int
				o, err = strconv.Atoi(buf.String())
				if err != nil {
					return
				}
				if o < 0 || 32 < o {
					err = fmt.Errorf("length out of range: %v", o)
					return
				}
				na.length = byte(o)
				return
			}
			err = nil
			return
		case err != nil:
			return
		case '0' <= b && b <= '9':
			buf.WriteByte(b)
			continue
		case b == '.' || b == '/':
			var o int
			o, err = strconv.Atoi(buf.String())
			if err != nil {
				return
			}
			if o < 0 || 255 < o {
				err = fmt.Errorf("octet out of range: %v", o)
				return
			}
			na.address[i] = byte(o)
			i++
			buf = new(bytes.Buffer)
			if b == '/' {
				onLen = true
			}
		}
	}
	err = fmt.Errorf("input is longer than expected")
	return
}

/*
IP parses a string as IPv4 or IPv6, returning a NetAddr or
an error.

Address strings without lengths are always interpreted as
single addresses with maximal length.

Sample inputs:

  - 192.0.2.9/24
  - 2001:DB8::/32
  - 192.0.2.1
  - 2001:db8::
  - 192.0.2/24
*/
func IP(s string) (na *NetAddr, err error) {
	var ip net.IP
	var ipNet *net.IPNet

	ip, ipNet, err = net.ParseCIDR(s)
	switch err.(type) {
	case nil:
	case *net.ParseError:
		ip = net.ParseIP(s)
		if ip == nil {
			na, err = parseTruncatedIPv4(s)
			if err != nil {
				err = fmt.Errorf("invalid ip address %v: %v", s, err)
			}
			return
		}
		if test := ip.To4(); test != nil {
			ipNet = &net.IPNet{
				IP:   ip,
				Mask: net.IPMask([]byte{255, 255, 255, 255}),
			}
		} else if test := ip.To16(); test != nil {
			ipNet = &net.IPNet{
				IP: ip,
				Mask: net.IPMask(
					[]byte{
						0xff, 0xff, 0xff, 0xff,
						0xff, 0xff, 0xff, 0xff,
						0xff, 0xff, 0xff, 0xff,
						0xff, 0xff, 0xff, 0xff,
					},
				),
			}
		}
		err = nil
	default:
		return
	}
	var length, _ int = net.IPMask.Size(ipNet.Mask)
	var address []byte = ip.To4()

	if address == nil {
		address = ip.To16()
	}
	if address == nil {
		panic(fmt.Sprintf("unexpected failure parsing '%s'\n", s))
	}
	na = &NetAddr{address, uint8(length)}
	return
}