src/Peer.hs

   1 module Peer where
   2
   3 import Prelude hiding (lookup, concat, replicate, splitAt)
   4
   5 import Bencode (BVal(..), InfoDict, decode)
   6 import Data.ByteString.Char8 (ByteString, pack, unpack, concat, replicate, splitAt)
   7 import Data.ByteString.Lazy (toChunks)
   8 import Data.Int (Int8)
   9 import Data.List (intercalate)
  10 import Data.Map as M ((!), lookup)
  11 import Tracker (infoHash)
  12 import Utils (splitN)
  13 import qualified Data.Binary as Bin (encode)
  14 import qualified Data.ByteString.Base16 as B16 (encode)
  15
  16
  17 type Address = String
  18 type Port = Integer
  19
  20 data Peer = Peer Address Port
  21             deriving (Show)
  22
  23 data PeerResp = PeerResponse { interval :: Maybe Integer
  24                              , peers :: [Peer]
  25                              , complete :: Maybe Integer
  26                              , incomplete :: Maybe Integer
  27                              } deriving (Show)
  28
  29 toInt :: String -> Integer
  30 toInt = read
  31
  32 getPeers :: PeerResp -> [Peer]
  33 getPeers = peers
  34
  35 getPeerResponse :: ByteString -> PeerResp
  36 getPeerResponse body = case decode body of
  37                         Right (Bdict peerM) ->
  38                           let (Just (Bint i)) = lookup (Bstr (pack "lookup")) peerM
  39                               (Bstr peersBS) = peerM ! Bstr (pack "peers")
  40                               pl = map (\peer -> let (ip', port') = splitAt 4 peer
  41                                                  in Peer (toIPNum ip') (toPortNum port'))
  42                                    (splitN 6 peersBS)
  43                           in PeerResponse { interval = Just i
  44                                           , peers = pl
  45                                           , complete = Nothing
  46                                           , incomplete = Nothing
  47                                           }
  48                           where toPortNum = read . ("0x" ++) . unpack . B16.encode
  49                                 toIPNum = intercalate "." .
  50                                           map (show . toInt . ("0x" ++) . unpack) .
  51                                           splitN 2 . B16.encode
  52
  53                         _ -> PeerResponse { interval = Nothing
  54                                           , peers = []
  55                                           , complete = Nothing
  56                                           , incomplete = Nothing
  57                                           }
  58
  59
  60 handShakeMsg :: InfoDict -> String -> ByteString
  61 handShakeMsg m peer_id = let pstrlen = concat $ toChunks $ Bin.encode (19 :: Int8)
  62                              pstr = pack "BitTorrent protocol"
  63                              reserved = replicate 8 '\0'
  64                              infoH = infoHash m
  65                              peerID = pack peer_id
  66                          in concat [pstrlen, pstr, reserved, infoH, peerID]