Tsuru.hs

{-# LANGUAGE OverloadedStrings #-}

import Data.Conduit.List (sinkNull, mapMaybe) 
import Numeric (showFFloat, showInt)
import System.Console.GetOpt.Simple 
import Data.Conduit
import Data.List.Split (splitPlacesBlanks)
import Data.List (isSuffixOf)
import Network.Pcap.Conduit (Packet, sourceOffline, sourceLive)
import Network.Pcap.Base (hdrSeconds, hdrUseconds) 
import qualified Data.Map.Strict as M
import qualified Data.Heap as H
import qualified Data.ByteString.Char8 as C
import Control.Monad.IO.Class (liftIO)

type Bid = Offer
type Ask = Offer

data Offer = Offer { quantity :: Int, price :: Int} deriving Eq
instance Show Offer where
  show o = showInt (quantity o) $ '@':(show $ price o)

data Quote = Quote { pktTime :: Double, 
		     acceptTime :: Int, 
		     issueCode :: String, 
		     b5 :: Bid, 
		     b4 :: Bid, 
		     b3 :: Bid, 
		     b2 :: Bid, 
		     b1 :: Bid, 
		     a5 :: Ask, 
		     a4 :: Ask, 
		     a3 :: Ask, 
		     a2 :: Ask, 
		     a1 :: Ask 
		   } deriving Eq
instance Show Quote where
  show q = foldl1 (\a b -> a . (showChar '\t') . b) ([(showFFloat Nothing $ pktTime q), (showInt $ acceptTime q), (showString $ issueCode q)] ++ (map shows [b5 q, b4 q, b3 q, b2 q, b1 q, a1 q, a2 q, a3 q, a4 q, a5 q])) ""
instance Ord Quote where
  compare q1 q2 = compare (acceptTime q1) (acceptTime q2)

-- Convert a single packet to quote, discarding if it is not a valid quote packet.
packetToQuote :: Packet -> Maybe Quote
packetToQuote (hdr, msg) = if C.null content then Nothing else Just quote 
	where (_,content) = C.breakSubstring "B6034" msg
	      [_,_,_,issCode,_,_,_,bp1,bq1,bp2,bq2,bp3,bq3,bp4,bq4,bp5,bq5,_,ap1,aq1,ap2,aq2,ap3,aq3,ap4,aq4,ap5,aq5,_,_,_,_,_,_,_,_,_,_,_,_,accTim,_] = splitPlacesBlanks [2,2,1,12,3,2,7,5,7,5,7,5,7,5,7,5,7,7,5,7,5,7,5,7,5,7,5,7,5,4,4,4,4,4,5,4,4,4,4,4,8,1] (C.unpack content) -- Splitting this into so many pieces is not really necessary in this case, but is good practice in case we ever want to refactor.
	      pTime = (fromIntegral $ hdrSeconds hdr) + (fromIntegral $ hdrUseconds hdr)/1000
	      quote = Quote { pktTime = pTime, acceptTime = read accTim, issueCode = issCode, b5 = mkOffer bp5 bq5, b4 = mkOffer bp4 bq4, b3 = mkOffer bp3 bq3, b2 = mkOffer bp2 bq2, b1 = mkOffer bp1 bq1, a5 = mkOffer ap5 aq5, a4 = mkOffer ap4 aq4, a3 = mkOffer ap3 aq3, a2 = mkOffer ap2 aq2, a1 = mkOffer ap1 aq1 }
	      mkOffer p q = Offer { quantity = read q, price = read p }


-- Reorder a stream of packets according to their quote accept times.
-- Maintains a min-heap cache of at most 3 seconds worth of packets at a time.
reorder :: Conduit Quote IO Quote
reorder = rloop H.empty
  where 
    yieldHeap h = do
      if H.size h == 1 
      then yield $ H.minimum h
      else do
        yield $ H.minimum h
        yieldHeap $ H.deleteMin h
    rloop h = do
      x <- await
      case x of 
	Nothing -> yieldHeap h
	Just q -> 
	  do 
	    let heap = H.insert q h
	    if acceptTime (H.minimum heap) + 300 < acceptTime q 
	      then 
	        do
		  yield $ H.minimum heap
		  rloop $ H.deleteMin heap
	      else rloop heap

-- Given a stream of quotes, print them.
printQuote :: Sink Quote IO ()
printQuote = do
  x <- await
  case x of 
    Nothing -> sinkNull
    Just y -> do 
      liftIO $ putStrLn (show y) 
      printQuote

-- Convert a packet stream to a stream of quotes, discarding any invalid packets,
-- optionally reorder quotes by exchange accept time, then print each quote in order.
printPackets :: Bool -> Conduit () IO Packet -> IO ()
printPackets False packets = packets $= mapMaybe packetToQuote $$ printQuote
printPackets True packets = packets $= mapMaybe packetToQuote $= reorder $$ printQuote

cliOptions :: Conf
cliOptions = [ (noArg, "reorder", Optional, "Reorder messages by quote accept time, within a 3-second margin."), 
	       (noArg, "live-capture", Optional, "Instead of parsing pcap files, capture packets directly. May require root privileges."), 
	       (arg, "interface", Default "any", "Specify an interface for live-capture. Defaults to 'any', i.e. capture on all available interfaces."), 
	       (arg, "time", Default "0", "Set timeout interval in microseconds for live packet capture. Defaults to 0, i.e. wait indefinitely."), 
	       (noArg, "matrixmode", Optional, "The machines tell elegant lies."), 
	       (arg, "size", Default "4096", "Set snapshot length (in bytes) for live packet capture. Defaults to 4096.")]

main :: IO ()
main = do
  (opts, args) <- getUsingConf cliOptions []
  let r = M.member "reorder" opts 
  if M.member "matrixmode" opts then putStr "\ESC[32m" else putStr ""
  if M.member "live-capture" opts 
    then printPackets r $ sourceLive (read $ opts M.! "interface") (read $ opts M.! "size") True (read $ opts M.! "time") 
    else 
      case filter (".pcap" `isSuffixOf`) args of
        [] -> putStrLn "You must either specify some .pcap files or do a live capture."
        files -> mapM_ (printPackets r . sourceOffline) files