Parsing pcap files

{- Karol Samborski <edv.karol@gmail.com> -}
module Main where

import System.Environment
import System.Time
import Control.Applicative
import Network.Pcap
import Data.List
import qualified Data.ByteString.Char8 as BS

data Options = Options { filename :: FilePath, quoteOrder :: Bool }
data AcceptTime = AcceptTime {
      hh :: BS.ByteString
    , mm :: BS.ByteString
    , ss :: BS.ByteString
    , uu :: BS.ByteString
    }

instance Show AcceptTime where
    show (AcceptTime h m s u) = (BS.unpack h) ++ ":"
                                              ++ (BS.unpack m)
                                              ++ ":"
                                              ++ (BS.unpack s)
                                              ++ "."
                                              ++ (BS.unpack u)

instance Eq AcceptTime where
    (==) a1 a2 = t1 == t2
      where
        t1 = BS.append (hh a1) $ BS.append (mm a1) $ BS.append (ss a1) (uu a1)
        t2 = BS.append (hh a2) $ BS.append (mm a2) $ BS.append (ss a2) (uu a2)

instance Ord AcceptTime where
    compare a1 a2 = if t1 > t2 then GT else if t1 < t2 then LT else EQ
      where
        t1 = BS.append (hh a1) $ BS.append (mm a1) $ BS.append (ss a1) (uu a1)
        t2 = BS.append (hh a2) $ BS.append (mm a2) $ BS.append (ss a2) (uu a2)


data MarketData = MarketData {
      issueCode  :: BS.ByteString
    , bids       :: [(BS.ByteString, BS.ByteString)]
    , asks       :: [(BS.ByteString, BS.ByteString)]
    , acceptTime :: AcceptTime
    }

instance Show MarketData where
    show (MarketData i b a t) = (show t)
        ++ " " ++ (BS.unpack i)
        ++ " " ++ (drop 1 $ foldr (\(ty, pr) ac ->
            "@" ++ (BS.unpack ty) ++ "@" ++ (BS.unpack pr) ++ ac) [] b)
        ++ " " ++ (drop 1 $ foldr (\(ty, pr) ac ->
            "@" ++ (BS.unpack ty) ++ "@" ++ (BS.unpack pr) ++ ac) [] a)

parseData :: BS.ByteString
          -> MarketData
parseData line = MarketData (BS.take 12 l) _bids _asks (AcceptTime h m s u)
  where
    l = BS.drop 5 line
    (_bids, r1) = takeBids $ BS.drop 24 l
    (_asks, r2) = takeAsks $ BS.drop 7 r1
    (h, r3) = BS.splitAt 2 (BS.take 8 $ BS.drop 50 r2)
    (m, r4) = BS.splitAt 2 r3
    (s, u) = BS.splitAt 2 r4

takeBids :: BS.ByteString
         -> ([(BS.ByteString, BS.ByteString)], BS.ByteString)
takeBids bs = takeBids' bs (5::Integer)
  where
    takeBids' b n
        | n > 0     = let (_b, _r) = takeBids' (BS.drop 12 b) (n-1)
                      in (_b ++ [(BS.splitAt 5 $ BS.take 12 b)], _r)
        | otherwise = ([], b)

takeAsks :: BS.ByteString
         -> ([(BS.ByteString, BS.ByteString)], BS.ByteString)
takeAsks bs = takeAsks' bs (5::Integer)
  where
    takeAsks' b n
        | n > 0     = let (_b, _r) = takeAsks' (BS.drop 12 b) (n-1)
                      in ((BS.splitAt 5 $ BS.take 12 b):_b, _r)
        | otherwise = ([], b)

prefix :: BS.ByteString
prefix = (BS.pack "B6034")

showData :: PktHdr
         -> BS.ByteString
         -> IO ()
showData hdr dat = do
    let (_, _d) = BS.breakSubstring prefix dat
        (sec, msec) = (hdrTime hdr) `divMod` 1000000
    if (BS.length _d) > 0
        then putStrLn $ (show $ TOD (toInteger sec) (toInteger (msec*1000000)))
            ++ " " ++ (show $ parseData _d)
        else return ()

getData :: PktHdr
        -> BS.ByteString
        -> (CalendarTime,MarketData)
getData hdr dat = (toUTCTime $ TOD (toInteger sec) (toInteger (msec*1000000)), parseData dat)
  where
    (sec, msec) = (hdrTime hdr) `divMod` 1000000

nextBS' :: PcapHandle
        -> IO (PktHdr,BS.ByteString)
nextBS' h = do
    (hdr, dat) <- nextBS h
    let (_, _d) = BS.breakSubstring prefix dat
    if (hdrCaptureLength hdr) == 0
        then return (hdr, dat)
        else if ((BS.length _d) > 0)
            then return (hdr,_d)
            else nextBS' h

compareData :: (CalendarTime,MarketData)
            -> (CalendarTime,MarketData)
            -> Ordering
compareData (_,m1) (_,m2) = if (a1 > a2)
    then GT
    else if (a1 < a2)
        then LT
        else EQ
  where
    a1 = acceptTime m1
    a2 = acceptTime m2

readSortQuote :: Int
              -> [(CalendarTime,MarketData)]
              -> PcapHandle
              -> IO ()
readSortQuote sec buf h = do
    (hdr, dat) <- nextBS' h
    let rec@(t,_dat) = getData hdr dat
    if (hdrCaptureLength hdr) == 0
        then mapM_ (\(_t, _m) ->
            putStrLn $ (calendarTimeToString _t) ++ " " ++ show _m)
                     $ sortBy compareData buf
        else if (ctSec t) > (sec + 3)
            then let (s,r) = break (\(time,_) -> (ctSec time) > sec)
                             $ sortBy compareData buf
                 in mapM_ (\(_t, _m) ->
                    putStrLn $ (calendarTimeToString _t) ++ " " ++ show _m) s
                    >> readSortQuote (ctSec $ fst $ head r) (r++[rec]) h
            else readSortQuote sec (buf++[rec]) h

readQuote :: Bool
          -> PcapHandle
          -> IO ()
readQuote False h = dispatchBS h (-1) showData >> return ()
readQuote True h = do
    (_hdr, _dat) <- nextBS' h
    let first@(time,_) = getData _hdr _dat
    readSortQuote (ctSec time) [first] h

main :: IO ()
main = do
    opt <- liftA parseOptions getArgs
    hdl <- openOffline $ filename opt
    readQuote (quoteOrder opt) hdl

parseOptions :: [String]
             -> Options
parseOptions [] = Options "" False
parseOptions (x:xs)
    | x == "-r" = Options (head xs) True
    | otherwise = Options x False