{-# LANGUAGE MultiParamTypeClasses, FunctionalDependencies, RecordWildCards #-}

{-|

This module provides @Get@ values for parsing various 
IP packets and headers from ByteStrings into a byte-sequence-independent 
representation as Haskell datatypes. 

Warning: 

These are incomplete. The headers may not contain all the information
that the protocols specify. For example, the Haskell representation of an IP Header
only includes source and destination addresses and IP protocol number, even though
an IP packet has many more header fields. More seriously, an IP header may have an optional 
extra headers section after the destination address. We assume this is not present. If it is present, 
then the transport protocol header will not be directly after the destination address, but will be after 
these options. Therefore functions that assume this, such as the getExactMatch function below, will give 
incorrect results when applied to such IP packets. 

The Haskell representations of the headers for the transport protocols are similarly incomplete. 
Again, the Get instances for the transport protocols may not parse through the end of the 
transport protocol header. 

-}
module Nettle.IPv4.IPPacket ( 
  -- * IP Packet 
  IPPacket(..)
  , IPHeader(..)
  , DifferentiatedServicesCodePoint
  , FragOffset
  , IPProtocol
  , IPTypeOfService
  , TransportPort
  , ipTypeTcp 
  , ipTypeUdp 
  , ipTypeIcmp
  , IPBody(..)
    
    -- * Parsers
  , getIPPacket
  , getIPHeader
  , ICMPHeader
  , ICMPType
  , ICMPCode
  , getICMPHeader
  , TCPHeader
  , TCPPortNumber
  , getTCPHeader
  , UDPHeader
  , UDPPortNumber
  , getUDPHeader
  
  -- * Framed Messages 
  , FramedMessage(..)
  
  ) where

import Nettle.IPv4.IPAddress
import Data.Bits
import Data.Word
import Data.Binary 
import Data.Binary.Get
import qualified Data.ByteString.Lazy as B

-- | An IP packet consists of a header and a body.
data IPPacket = IPPacket IPHeader IPBody
                deriving (Show,Eq)

-- | An IP Header includes various information about the packet, including the type of payload it contains. 
-- Warning: this definition does not include every header field included in an IP packet. 
data IPHeader = IPHeader { ipSrcAddress :: IPAddress
                         , ipDstAddress :: IPAddress
                         , ipProtocol   :: IPProtocol  
                         , headerLength :: Int
                         , totalLength :: Int
                         , dscp         :: DifferentiatedServicesCodePoint -- ^ differentiated services code point - 6 bit number
                         }
                deriving (Read,Show,Eq)

type DifferentiatedServicesCodePoint = Word8
type FragOffset      = Word16
type IPProtocol      = Word8
type IPTypeOfService = Word8
type TransportPort   = Word16

ipTypeTcp, ipTypeUdp, ipTypeIcmp :: IPProtocol

ipTypeTcp  = 6
ipTypeUdp  = 17
ipTypeIcmp = 1

-- | The body of an IP packet can be either a TCP, UDP, ICMP or other packet. 
-- Packets other than TCP, UDP, ICMP are represented as unparsed @ByteString@ values.
data IPBody = TCPInIP TCPHeader  
            | UDPInIP UDPHeader  
            | ICMPInIP ICMPHeader
            | UninterpretedIPBody B.ByteString
              deriving (Show,Eq)

-- | The @FramedMessage@ class defines the class of data types
-- which represents messages having source and destination addresses
-- and a message body.
class FramedMessage m a b | m -> a, m -> b where
    body          :: m -> b
    sourceAddress :: m -> a
    destAddress   :: m -> a

instance FramedMessage IPPacket IPAddress IPBody where
    body (IPPacket _ b) = b
    sourceAddress (IPPacket hdr _) = ipSrcAddress hdr
    destAddress (IPPacket hdr _)   = ipDstAddress hdr


getIPHeader :: Get IPHeader
getIPHeader = do 
  b1                 <- getWord8
  diffServ           <- getWord8
  totalLen           <- getWord16be
  ident              <- getWord16be
  flagsAndFragOffset <- getWord16be
  ttl                <- getWord8
  nwproto            <- getIPProtocol
  hdrChecksum        <- getWord16be
  nwsrc              <- getIPAddress
  nwdst              <- getIPAddress
  return (IPHeader { ipSrcAddress = nwsrc 
                   , ipDstAddress = nwdst 
                   , ipProtocol = nwproto
                   , headerLength = fromIntegral (b1 .&. 0x0f)
                   , totalLength  = fromIntegral totalLen
                   , dscp = shiftR diffServ 2
                   } )

getIPProtocol :: Get IPProtocol 
getIPProtocol = getWord8

getFragOffset :: Get FragOffset
getFragOffset = getWord16be

getIPPacket :: Get IPPacket
getIPPacket = do 
  hdr  <- getIPHeader
  body <- getIPBody hdr
  return (IPPacket hdr body)
    where getIPBody (IPHeader {..}) 
              | ipProtocol == ipTypeTcp  = getTCPHeader  >>= return . TCPInIP
              | ipProtocol == ipTypeUdp  = getUDPHeader  >>= return . UDPInIP
              | ipProtocol == ipTypeIcmp = getICMPHeader >>= return . ICMPInIP
              | otherwise             = getLazyByteString (fromIntegral (totalLength - headerLength)) >>= 
                                        return . UninterpretedIPBody

-- Transport Header

type ICMPHeader = (ICMPType, ICMPCode)
type ICMPType = Word8
type ICMPCode = Word8

getICMPHeader :: Get ICMPHeader
getICMPHeader = do 
  icmp_type <- getWord8
  icmp_code <- getWord8
  skip 6
  return (icmp_type, icmp_code)

type TCPHeader  = (TCPPortNumber, TCPPortNumber)
type TCPPortNumber = Word16

getTCPHeader :: Get TCPHeader
getTCPHeader = do 
  srcp <- getWord16be
  dstp <- getWord16be
  return (srcp,dstp)

type UDPHeader     = (UDPPortNumber, UDPPortNumber)
type UDPPortNumber = Word16

getUDPHeader :: Get UDPHeader
getUDPHeader = do 
  srcp <- getWord16be
  dstp <- getWord16be
  return (srcp,dstp)