{-# LANGUAGE MultiParamTypeClasses #-}

-- | This module provides data structures for Ethernet frames
-- as well as parsers and unparsers for Ethernet frames. 
module Nettle.Ethernet.EthernetFrame ( 
  
  -- * Data types
  EthernetFrame(..)
  , EthernetHeader(..)
  , EthernetTypeCode  
  , ethTypeVLAN
  , ethTypeIP
  , ethTypeARP
  , ethTypeLLDP
  , typeEth2Cutoff
  , VLANPriority
  , VLANID
  , EthernetBody(..)
  , ARPPacket(..)
  , ARPOpCode(..)

    -- * Parsers and unparsers 
  , GetE
  , ErrorMessage
  , runGetE
  , getEthernetFrame  
  , getEthHeader
  , putEthHeader
  , getARPPacket        
  ) where

import Nettle.Ethernet.EthernetAddress
import Nettle.IPv4.IPPacket
import Nettle.IPv4.IPAddress
import qualified Data.ByteString.Lazy as B
import Data.Binary
import Data.Binary.Get
import Data.Binary.Put
import Data.Word
import Data.Bits
import Control.Monad
import Control.Monad.Error

-- It was based on http://en.wikipedia.org/wiki/File:Ethernet_Type_II_Frame_format.svg

-- | An Ethernet frame consists of an Ethernet header and an Ethernet body.
data EthernetFrame = EthernetFrame EthernetHeader EthernetBody 
                     deriving (Show,Eq)

data EthernetHeader   = EthernetHeader { destMACAddress   :: EthernetAddress, 
                                         sourceMACAddress :: EthernetAddress, 
                                         typeCode         :: EthernetTypeCode }
                      | Ethernet8021Q {  destMACAddress           :: EthernetAddress, 
                                         sourceMACAddress         :: EthernetAddress, 
                                         typeCode                 :: EthernetTypeCode, 
                                         priorityCodePoint        :: VLANPriority, 
                                         canonicalFormatIndicator :: Bool, 
                                         vlanId                   :: VLANID }
                        deriving (Read,Show,Eq)

type VLANPriority     = Word8

-- | Ethernet type code, determines the type of payload carried by an Ethernet frame.
type EthernetTypeCode = Word16

type VLANID           = Word16

-- | The body of an Ethernet frame is either an IP packet, an ARP packet, or an uninterpreted @ByteString@
data EthernetBody  = IPInEthernet IPPacket 
                   | ARPInEthernet ARPPacket
                   | UninterpretedEthernetBody B.ByteString
                     deriving (Show,Eq)

instance FramedMessage EthernetFrame EthernetAddress EthernetBody where
    body          (EthernetFrame _   b) = b
    sourceAddress (EthernetFrame hdr _) = sourceMACAddress hdr
    destAddress   (EthernetFrame hdr _) = destMACAddress hdr

-- | An ARP packet
data ARPPacket = ARPPacket { arpOpCode             :: ARPOpCode
                           , senderEthernetAddress :: EthernetAddress
                           , senderIPAddress       :: IPAddress
                           , targetEthernetAddress :: EthernetAddress
                           , targetIPAddress       :: IPAddress
                           } deriving (Show,Eq)

-- | Type of ARP message.
data ARPOpCode = ARPRequest 
               | ARPReply deriving (Show,Eq)                             

-- | Type of parsers that can fail with an error message
type GetE a = ErrorT ErrorMessage Get a

-- | When a @GetE@ parser fails, it provides an error message as a string
type ErrorMessage = String

-- | Method to run a @GetE@ parser
runGetE :: GetE a -> B.ByteString -> Either ErrorMessage a
runGetE g = runGet (runErrorT g)

-- | Parser for Ethernet frames.
getEthernetFrame :: GetE EthernetFrame
getEthernetFrame = do 
  hdr <- getEthHeader
  if typeCode hdr == ethTypeIP
    then do ipPacket <- lift getIPPacket
            return $ EthernetFrame hdr (IPInEthernet ipPacket)
    else if typeCode hdr == ethTypeARP 
         then do arpPacket <- getARPPacket
                 return (EthernetFrame hdr (ARPInEthernet arpPacket))
         else do body <- lift getRemainingLazyByteString
                 return $ EthernetFrame hdr (UninterpretedEthernetBody body)

-- | Parser for Ethernet headers.
getEthHeader :: GetE EthernetHeader
getEthHeader = do 
  dstAddr <- lift getEthernetAddress
  srcAddr <- lift getEthernetAddress
  tcode   <- lift getWord16be
  if tcode < typeEth2Cutoff 
    then throwError ("tcode of ethernet header is not greater than " ++ show typeEth2Cutoff)
    else if (tcode == ethTypeVLAN) 
         then do x <- lift getWord16be
                 etherType <- lift getWord16be
                 let pcp = fromIntegral (shiftR x 13)
                 let cfi = testBit x 12
                 let vid = clearBits x [12,13,14,15]
                 return (Ethernet8021Q dstAddr srcAddr etherType pcp cfi vid)
         else return (EthernetHeader dstAddr srcAddr tcode)

-- | Unparser for Ethernet headers.
putEthHeader :: EthernetHeader -> Put 
putEthHeader (EthernetHeader dstAddr srcAddr tcode) =  
    do putEthernetAddress dstAddr
       putEthernetAddress srcAddr
       putWord16be tcode
putEthHeader (Ethernet8021Q dstAddr srcAddr tcode pcp cfi vid) = 
    do putEthernetAddress dstAddr
       putEthernetAddress srcAddr
       putWord16be ethTypeVLAN
       putWord16be x
       putWord16be tcode
    where x = let y = shiftL (fromIntegral pcp :: Word16) 13
                  y' = if cfi then setBit y 12 else y
              in y' + fromIntegral vid



ethTypeIP, ethTypeARP, ethTypeLLDP, ethTypeVLAN, typeEth2Cutoff :: EthernetTypeCode
ethTypeIP           = 0x0800
ethTypeARP          = 0x0806
ethTypeLLDP         = 0x88CC
ethTypeVLAN         = 0x8100
typeEth2Cutoff = 0x0600


-- | Parser for ARP packets
getARPPacket :: GetE ARPPacket
getARPPacket = do 
  htype <- lift getWord16be
  ptype <- lift getWord16be
  hlen  <- lift getWord8
  plen  <- lift getWord8
  opCode <- getARPOpCode
  sha <- lift getEthernetAddress
  spa <- lift getIPAddress
  tha <- lift getEthernetAddress
  tpa <- lift getIPAddress
  return (ARPPacket { arpOpCode             = opCode, 
                      senderEthernetAddress = sha, 
                      senderIPAddress       = spa,
                      targetEthernetAddress = tha, 
                      targetIPAddress       = tpa })
  
getARPOpCode :: GetE ARPOpCode  
getARPOpCode = do 
  op <- lift getWord16be
  if op == 1 
    then return ARPRequest 
    else if op == 2 
         then return ARPReply 
         else throwError ("Unrecognized ARP opcode: " ++ show op)

clearBits :: Bits a => a -> [Int] -> a 
clearBits = foldl clearBit