{-# LANGUAGE DisambiguateRecordFields, FlexibleInstances, MultiParamTypeClasses, DeriveDataTypeable #-}
{- |The Data.IP library exports IPv4 address and header structures.

   FIXME:
   There is currently no support for options fields of the IP header.
 -}
module Data.IP
	( IPv4(..)
	, IPv4Header(..)
	, IPv4Flag(..)
	, IP
	, IPHeader
	, dummyIPv4Header
	, module Data.IPv6
	) where

import Control.Monad (sequence, when, liftM)
import qualified Data.ByteString as B
import Data.Serialize
import Data.Serialize.Put
import Data.Serialize.Get
import Data.CSum
import Data.Data
import Data.List
import Data.IPv6
import Data.Header
import Data.Bits
import Text.PrettyPrint
import Text.PrettyPrint.HughesPJClass
import Data.Word

type IP = Either IPv4 IPv6

type IPHeader = Either IPv4Header IPv6Header

-- |For IPv4 addresses.
data IPv4 = IPv4 Word32 deriving (Eq, Ord, Show, Read, Data, Typeable)

instance Serialize IPv4 where
	put (IPv4 b) = putWord32be b
	get = liftM IPv4 getWord32be

-- |Don't fragment, more fragment and reserved flags
data IPv4Flag = DF | MF | Res deriving (Eq, Ord, Show, Read, Data, Typeable)

instance Enum [IPv4Flag] where
	fromEnum xs = foldl' (.|.) 0 $ map fromEnum1 xs
	toEnum f = map snd $ filter fst [(testBit f 0, Res), (testBit f 1, MF), (testBit f 2, DF)]

fromEnum1 DF   = 4
fromEnum1 MF   = 2
fromEnum1 Res  = 1

-- |This IPv4 header structure lacks support for options.  Ints are used
-- for most integral data types and the binary instance hands the bit packing.
--
-- No warning is provided if a value is trunkated when packed!
data IPv4Header =
	IPv4Hdr { hdrLength		:: Int
		, version		:: Int
		, tos			:: Int
		, totalLength		:: Int
		, ipID			:: Int
		, flags			:: [IPv4Flag]
		, fragmentOffset	:: Int
		, ttl			:: Int
		, protocol		:: Int
		, checksum		:: CSum
		, source		:: IPv4
		, destination		:: IPv4
	} deriving (Eq, Ord, Show, Read, Data, Typeable)

-- |A dummy header with zeroed fields except version, header length and TTL (255).
dummyIPv4Header = IPv4Hdr 5 4 0 0 0 [] 0 255 0 0 ipv4zero ipv4zero

ipv4zero = IPv4 0

instance Serialize IPv4Header where
  put (IPv4Hdr ihl ver tos len id flags off ttl prot csum src dst) = do
	pW8 $ (ihl .&. 0xF) .|. (ver `shiftL` 4 .&. 0xF0)
	pW8 tos
	pW16 len
	pW16 id
	let offFlags = (off .&. 0x1FFF) .|. fromIntegral (fromEnum flags `shiftL` 13)
	pW16 offFlags
	pW8 ttl
	pW8 prot
	put csum
	put src
	put dst

  get = do
	ihlVer <- gW8
	let ihl = (ihlVer .&. 0xF)
	    ver = (ihlVer `shiftR` 4) .&. 0xF
	tos <- gW8
	len <- gW16
	id  <- gW16
	offFlags <- gW16
	let off = offFlags .&. 0x1FFF
	    flags = toEnum $ offFlags `shiftR` 13
	ttl <- gW8
	prot <- gW8
	csum <- get
	src <- get
	dst <- get
	return $ IPv4Hdr ihl ver tos len id flags off ttl prot csum src dst

gW8 = getWord8 >>= return . fromIntegral
gW16 = getWord16be >>= return . fromIntegral
pW8 = putWord8 . fromIntegral
pW16 = putWord16be . fromIntegral
pW32 = putWord32be . fromIntegral

-- L3Header and L3Address instances (see Data.Header)
instance L3Header IPv4Header IPv4 CSum where
	getChecksum = checksum
	setChecksum h c = h { checksum = c }
	src = source
	dst = destination
	pseudoHeader h = runPut (do
		put (src h)
		put (dst h)
		putWord8 0
		pW8 $ fromIntegral (protocol h)
		pW16 (totalLength h))
	computeChecksum h = csum16 (encode (zeroChecksum h))

instance L3Address IPv4 IPv4Header where
	localBroadcast (IPv4 a) = IPv4 (0xFFFFFF00 .|. (0x000000FF .&. a))
	globalBroadcast = IPv4 0xFFFFFFFF

-- Pretty Printing and parsing instances
instance Pretty IPv4 where
	pPrint (IPv4 i) = text . concat . intersperse "." . map show $ unfoldr (\(i,v) -> if i /= 0 then Just (v .&. 0xFF, (i-1, v `shiftR` 8)) else Nothing) (4,i)