+import Prelude hiding (lookup, concat, replicate, splitAt, take, drop)
+
+import Control.Monad.State
+import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, drop, empty)
+import Data.Bits
+import Data.Word (Word8)
+import Data.Map ((!), adjust)
+import Network (connectTo, PortID(..))
+import System.IO (Handle, BufferMode(..), hSetBuffering, hClose)
+
+import FuncTorrent.Metainfo (Metainfo(..))
+import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg)
+import FuncTorrent.Utils (splitNum, verifyHash)
+import FuncTorrent.PieceManager (PieceDlState(..), PieceData(..), PieceMap, pickPiece, updatePieceAvailability)
+import qualified FuncTorrent.FileSystem as FS (MsgChannel, writePieceToDisk)
+
+data PState = PState { handle :: Handle
+ , peer :: Peer
+ , meChoking :: Bool
+ , meInterested :: Bool
+ , heChoking :: Bool
+ , heInterested :: Bool}
+
+havePiece :: PieceMap -> Integer -> Bool
+havePiece pm index =
+ dlstate (pm ! index) == Have
+
+connectToPeer :: Peer -> IO Handle
+connectToPeer (Peer ip port) = do
+ h <- connectTo ip (PortNumber (fromIntegral port))
+ hSetBuffering h LineBuffering
+ return h
+
+doHandshake :: Bool -> Handle -> Peer -> ByteString -> String -> IO ()
+doHandshake True h p infohash peerid = do
+ let hs = genHandshakeMsg infohash peerid
+ hPut h hs
+ putStrLn $ "--> handhake to peer: " ++ show p
+ _ <- hGet h (length (unpack hs))
+ putStrLn $ "<-- handshake from peer: " ++ show p
+ return ()
+doHandshake False h p infohash peerid = do
+ let hs = genHandshakeMsg infohash peerid
+ putStrLn "waiting for a handshake"
+ hsMsg <- hGet h (length (unpack hs))
+ putStrLn $ "<-- handshake from peer: " ++ show p
+ let rxInfoHash = take 20 $ drop 28 hsMsg
+ if rxInfoHash /= infohash
+ then do
+ putStrLn "infoHashes does not match"
+ hClose h
+ return ()
+ else do
+ _ <- hPut h hs
+ putStrLn $ "--> handhake to peer: " ++ show p
+ return ()
+
+bitfieldToList :: [Word8] -> [Integer]
+bitfieldToList bs = go bs 0
+ where go [] _ = []
+ go (b:bs') pos =
+ let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
+ in
+ setBits ++ go bs' (pos + 1)
+
+-- helper functions to manipulate PeerState
+toPeerState :: Handle
+ -> Peer
+ -> Bool -- ^ meChoking
+ -> Bool -- ^ meInterested
+ -> Bool -- ^ heChoking
+ -> Bool -- ^ heInterested
+ -> PState
+toPeerState h p meCh meIn heCh heIn =
+ PState { handle = h
+ , peer = p
+ , heChoking = heCh
+ , heInterested = heIn
+ , meChoking = meCh
+ , meInterested = meIn }
+
+handlePeerMsgs :: Peer -> String -> Metainfo -> PieceMap -> Bool -> FS.MsgChannel -> IO ()
+handlePeerMsgs p peerId m pieceMap isClient c = do
+ h <- connectToPeer p
+ doHandshake isClient h p (infoHash m) peerId
+ let pstate = toPeerState h p False False True True
+ _ <- runStateT (msgLoop pieceMap c) pstate
+ return ()