module Cidr
( Cidr(..),
  cidr_from_string,
  cidr_tests,
  combine_all,
  contains,
  contains_proper,
  max_octet1,
  max_octet2,
  max_octet3,
  max_octet4,
  min_octet1,
  min_octet2,
  min_octet3,
  min_octet4,
  prop_all_cidrs_contain_themselves,
  prop_contains_proper_intransitive
) where

import Data.List (nubBy)
import Test.HUnit
import Test.QuickCheck

import qualified Bit as B
import IPv4Address
import ListUtils
import Maskable
import Maskbits
import Octet


data Cidr = None | Cidr { ipv4address :: IPv4Address,
                          maskbits :: Maskbits }
            deriving (Eq)


instance Show Cidr where
    show Cidr.None = "None"
    show cidr = (show (ipv4address cidr)) ++ "/" ++ (show (maskbits cidr))


instance Arbitrary Cidr where
    arbitrary = do
      ipv4 <- arbitrary :: Gen IPv4Address
      mask <- arbitrary :: Gen Maskbits
      return (Cidr ipv4 mask)

    coarbitrary _ = variant 0


-- Two CIDR ranges are equivalent if they have the same network bits
-- and the masks are the same.
equivalent :: Cidr -> Cidr -> Bool
equivalent Cidr.None Cidr.None = True
equivalent Cidr.None _ = False
equivalent _ Cidr.None = False
equivalent (Cidr addr1 mbits1) (Cidr addr2 mbits2) =
    (mbits1 == mbits2) && ((apply_mask addr1 mbits1 B.Zero) == (apply_mask addr2 mbits2 B.Zero))

-- Returns the mask portion of a CIDR address. That is, everything
-- after the trailing slash.
maskbits_from_cidr_string :: String -> Maskbits
maskbits_from_cidr_string s
    | length partlist == 2 = maskbits_from_string (partlist !! 1)
    | otherwise = Maskbits.None
      where
        partlist = (splitWith (`elem` "/") s)


-- Takes an IP address String in CIDR notation, and returns a list of
-- its octets (as Ints).
octets_from_cidr_string :: String -> [Octet]
octets_from_cidr_string s =
    map octet_from_string (take 4 (splitWith (`elem` "./") s))


cidr_from_string :: String -> Cidr
cidr_from_string s
    | addr  == IPv4Address.None = Cidr.None
    | mbits == Maskbits.None    = Cidr.None
    | otherwise = Cidr addr mbits
               where
                 addr = ipv4address_from_octets (oct1) (oct2) (oct3) (oct4)
                 oct1 = (octs !! 0)
                 oct2 = (octs !! 1)
                 oct3 = (octs !! 2)
                 oct4 = (octs !! 3)
                 octs  = octets_from_cidr_string s
                 mbits = maskbits_from_cidr_string s



min_host :: Cidr -> IPv4Address
min_host Cidr.None                 = IPv4Address.None
min_host (Cidr IPv4Address.None _) = IPv4Address.None
min_host (Cidr _ Maskbits.None)    = IPv4Address.None
min_host (Cidr addr mask) = apply_mask addr mask B.Zero


max_host :: Cidr -> IPv4Address
max_host Cidr.None                 = IPv4Address.None
max_host (Cidr IPv4Address.None _) = IPv4Address.None
max_host (Cidr _ Maskbits.None)    = IPv4Address.None
max_host (Cidr addr mask) = apply_mask addr mask B.One


min_octet1 :: Cidr -> Octet
min_octet1 cidr = octet1 (min_host cidr)

min_octet2 :: Cidr -> Octet
min_octet2 cidr = octet2 (min_host cidr)

min_octet3 :: Cidr -> Octet
min_octet3 cidr = octet3 (min_host cidr)

min_octet4 :: Cidr -> Octet
min_octet4 cidr = octet4 (min_host cidr)

max_octet1 :: Cidr -> Octet
max_octet1 cidr = octet1 (max_host cidr)

max_octet2 :: Cidr -> Octet
max_octet2 cidr = octet2 (max_host cidr)

max_octet3 :: Cidr -> Octet
max_octet3 cidr = octet3 (max_host cidr)

max_octet4 :: Cidr -> Octet
max_octet4 cidr = octet4 (max_host cidr)



-- Return true if the first argument (a CIDR range) contains the
-- second (another CIDR range). There are a lot of ways we can be fed
-- junk here. For lack of a better alternative, just return False when
-- we are given nonsense.
contains :: Cidr -> Cidr -> Bool
contains Cidr.None _ = False
contains _ Cidr.None = False
contains (Cidr _ Maskbits.None) _    = False
contains (Cidr IPv4Address.None _) _ = False
contains _ (Cidr _ Maskbits.None)    = False
contains _ (Cidr IPv4Address.None _) = False

-- If the number of bits in the network part of the first address is
-- larger than the number of bits in the second, there is no way that
-- the first range can contain the second. For, if the number of
-- network bits is larger, then the number of host bits must be
-- smaller, and if cidr1 has fewer hosts than cidr2, cidr1 most
-- certainly does not contain cidr2.
--
-- On the other hand, if the first argument (cidr1) has fewer (or the
-- same number of) network bits as the second, it can contain the
-- second. In this case, we need to check that every host in cidr2 is
-- contained in cidr1. If a host in cidr2 is contained in cidr1, then
-- at least mbits1 of an address in cidr2 will match cidr1. For
-- example,
--
-- cidr1 = 192.168.1.0/23, cidr2 = 192.168.1.100/24
--
-- Here, cidr2 contains all of 192.168.1.0 through
-- 192.168.1.255. However, cidr1 contains BOTH 192.168.0.0 through
-- 192.168.0.255 and 192.168.1.0 through 192.168.1.255. In essence,
-- what we want to check is that cidr2 "begins with" something that
-- cidr1 CAN begin with. Since cidr1 can begin with 192.168.1, and
-- cidr2 DOES, cidr1 contains cidr2..
--
-- The way that we check this is to apply cidr1's mask to cidr2's
-- address and see if the result is the same as cidr1's mask applied
-- to cidr1's address.
--
contains (Cidr addr1 mbits1) (Cidr addr2 mbits2)
    | mbits1 > mbits2 = False
    | otherwise = addr1masked == addr2masked
    where
      addr1masked = apply_mask addr1 mbits1 B.Zero
      addr2masked = apply_mask addr2 mbits1 B.Zero


contains_proper :: Cidr -> Cidr -> Bool
contains_proper cidr1 cidr2 =
    (cidr1 `contains` cidr2) && (not (cidr2 `contains` cidr1))


-- A CIDR range is redundant (with respect to the given list) if
-- another CIDR range in that list properly contains it.
redundant :: [Cidr] -> Cidr -> Bool
redundant cidrlist cidr = any ((flip contains_proper) cidr) cidrlist


-- First, we look at all possible pairs of cidrs, and combine the
-- adjacent ones in to a new list. Then, we concatenate that list with
-- the original one, and filter out all of the redundancies. If two
-- adjacent Cidrs are combined into a larger one, they will be removed
-- in the second step since the larger Cidr must contain the smaller
-- two.
combine_all :: [Cidr] -> [Cidr]
combine_all cidrs =
    combine_contained unique_cidrs
    where
      unique_cidrs = nubBy equivalent valid_cidr_combinations
      valid_cidr_combinations = filter (/= Cidr.None) cidr_combinations
      cidr_combinations =
          cidrs ++ [ (combine_adjacent x y) | x <- cidrs, y <- cidrs ]


-- Take a list of CIDR ranges and filter out all of the ones that are
-- contained entirelt within some other range in the list.
combine_contained :: [Cidr] -> [Cidr]
combine_contained cidrs =
    filter (not . (redundant cidrs)) cidrs


-- If the two Cidrs are not adjacent, return Cidr.None. Otherwise,
-- decrement the maskbits of cidr1 and return that; it will contain
-- both cidr1 and cidr2.
combine_adjacent :: Cidr -> Cidr -> Cidr
combine_adjacent cidr1 cidr2
  | not (adjacent cidr1 cidr2) = Cidr.None
  | (maskbits cidr1 == Zero) = Cidr.None
  | otherwise = cidr1 { maskbits = decrement (maskbits cidr1) }



-- Determine whether or not two CIDR ranges are adjacent. If two
-- ranges lie consecutively within the IP space, they can be
-- combined. For example, 10.1.0.0/24 and 10.0.1.0/24 are adjacent,
-- and can be combined in to 10.1.0.0/23.
adjacent :: Cidr -> Cidr -> Bool
adjacent Cidr.None _ = False
adjacent _ Cidr.None = False
adjacent cidr1 cidr2
  | mbits1 /= mbits2 = False
  | mbits1 == Maskbits.Zero = False -- They're equal.
  | otherwise = (mbits1 == (most_sig_bit_different addr1 addr2))
  where
    addr1 = ipv4address cidr1
    addr2 = ipv4address cidr2
    mbits1 = maskbits cidr1
    mbits2 = maskbits cidr2





-- HUnit Tests

test_min_host1 :: Test
test_min_host1 =
    TestCase $ assertEqual "The minimum host in 10.0.0.0/24 is 10.0.0.0" expected actual
    where
      actual = show $ min_host (cidr_from_string "10.0.0.0/24")
      expected = "10.0.0.0"


test_max_host1 :: Test
test_max_host1 =
    TestCase $ assertEqual "The maximum host in 10.0.0.0/24 is 10.0.0.255" expected actual
    where
      actual = show $ max_host (cidr_from_string "10.0.0.0/24")
      expected = "10.0.0.255"


test_equality1 :: Test
test_equality1 =
    TestCase $ assertEqual "10.1.1.0/23 equals itself" True (cidr1 == cidr1)
    where
      cidr1 = cidr_from_string "10.1.1.0/23"


test_contains1 :: Test
test_contains1 =
    TestCase $ assertEqual "10.1.1.0/23 contains 10.1.1.0/24" True (cidr1 `contains` cidr2)
    where
      cidr1 = cidr_from_string "10.1.1.0/23"
      cidr2 = cidr_from_string "10.1.1.0/24"


test_contains2 :: Test
test_contains2 =
    TestCase $ assertEqual "10.1.1.0/23 contains itself" True (cidr1 `contains` cidr1)
    where
      cidr1 = cidr_from_string "10.1.1.0/23"


test_contains_proper1 :: Test
test_contains_proper1 =
    TestCase $ assertEqual "10.1.1.0/23 contains 10.1.1.0/24 properly" True (cidr1 `contains_proper` cidr2)
    where
      cidr1 = cidr_from_string "10.1.1.0/23"
      cidr2 = cidr_from_string "10.1.1.0/24"


test_contains_proper2 :: Test
test_contains_proper2 =
    TestCase $ assertEqual "10.1.1.0/23 does not contain itself properly" False (cidr1 `contains_proper` cidr1)
    where
      cidr1 = cidr_from_string "10.1.1.0/23"


test_adjacent1 :: Test
test_adjacent1 =
    TestCase $ assertEqual "10.1.0.0/24 is adjacent to 10.1.1.0/24" True (cidr1 `adjacent` cidr2)
    where
      cidr1 = cidr_from_string "10.1.0.0/24"
      cidr2 = cidr_from_string "10.1.1.0/24"


test_adjacent2 :: Test
test_adjacent2 =
    TestCase $ assertEqual "10.1.0.0/23 is not adjacent to 10.1.0.0/24" False (cidr1 `adjacent` cidr2)
    where
      cidr1 = cidr_from_string "10.1.0.0/23"
      cidr2 = cidr_from_string "10.1.0.0/24"


test_adjacent3 :: Test
test_adjacent3 =
    TestCase $ assertEqual "10.1.0.0/24 is not adjacent to 10.2.5.0/24" False (cidr1 `adjacent` cidr2)
    where
      cidr1 = cidr_from_string "10.1.0.0/24"
      cidr2 = cidr_from_string "10.2.5.0/24"


test_adjacent4 :: Test
test_adjacent4 =
    TestCase $ assertEqual "10.1.1.0/24 is not adjacent to 10.1.2.0/24" False (cidr1 `adjacent` cidr2)
    where
      cidr1 = cidr_from_string "10.1.1.0/24"
      cidr2 = cidr_from_string "10.1.2.0/24"


test_combine_contained1 :: Test
test_combine_contained1 =
    TestCase $ assertEqual "10.0.0.0/8, 10.1.0.0/16, and 10.1.1.0/24 combine to 10.0.0.0/8" expected_cidrs (combine_contained test_cidrs)
    where
      cidr1 = cidr_from_string "10.0.0.0/8"
      cidr2 = cidr_from_string "10.1.0.0/16"
      cidr3 = cidr_from_string "10.1.1.0/24"
      expected_cidrs = [cidr1]
      test_cidrs = [cidr1, cidr2, cidr3]


test_combine_contained2 :: Test
test_combine_contained2 =
    TestCase $ assertEqual "192.168.3.0/23 does not contain 192.168.1.0/24" [cidr1, cidr2] (combine_contained [cidr1, cidr2])
    where
      cidr1 = cidr_from_string "192.168.3.0/23"
      cidr2 = cidr_from_string "192.168.1.0/24"


test_combine_all1 :: Test
test_combine_all1 =
    TestCase $ assertEqual "10.0.0.0/24 is adjacent to 10.0.1.0/24 and 10.0.3.0/23 contains 10.0.2.0/24" expected_cidrs (combine_all test_cidrs)
    where
      cidr1 = cidr_from_string "10.0.0.0/24"
      cidr2 = cidr_from_string "10.0.1.0/24"
      cidr3 = cidr_from_string "10.0.2.0/24"
      cidr4 = cidr_from_string "10.0.3.0/23"
      cidr5 = cidr_from_string "10.0.0.0/23"
      expected_cidrs = [cidr4, cidr5]
      test_cidrs = [cidr1, cidr2, cidr3, cidr4]


test_combine_all2 :: Test
test_combine_all2 =
    TestCase $ assertEqual "127.0.0.1/32 combines with itself recursively" expected_cidrs (combine_all test_cidrs)
    where
      cidr1 = cidr_from_string "127.0.0.1/32"
      expected_cidrs = [cidr1]
      test_cidrs = [cidr1, cidr1, cidr1, cidr1, cidr1]


cidr_tests :: [Test]
cidr_tests = [ test_min_host1,
               test_max_host1,
               test_equality1,
               test_contains1,
               test_contains2,
               test_contains_proper1,
               test_contains_proper2,
               test_adjacent1,
               test_adjacent2,
               test_adjacent3,
               test_adjacent4,
               test_combine_contained1,
               test_combine_contained2,
               test_combine_all1,
               test_combine_all2
             ]


-- QuickCheck Tests
prop_all_cidrs_contain_themselves :: Cidr -> Bool
prop_all_cidrs_contain_themselves cidr1 = cidr1 `contains` cidr1


-- If cidr1 properly contains cidr2, then by definition cidr2
-- does not properly contain cidr1.
prop_contains_proper_intransitive :: Cidr -> Cidr -> Property
prop_contains_proper_intransitive cidr1 cidr2 =
    (cidr1 `contains_proper` cidr2) ==>
      (not (cidr2 `contains_proper` cidr1))