Bencoding.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
module Bencoding where

import Control.Applicative
import Crypto.Hash.SHA1
import Data.Attoparsec.ByteString.Char8 (Parser)
import qualified Data.Attoparsec.ByteString.Char8 as P
import Data.ByteString as B
import Data.ByteString.Char8 as BC
import Data.Functor
import Data.Foldable as Foldable
import Data.List
import Data.Monoid
import Data.Word
import Lens.Family2
import Prelude hiding (take)

data BValue = String ByteString
            | Number Integer
            | List [BValue]
            | Dictionary [(ByteString, BValue)]
              deriving(Eq,Show)

string :: Parser BValue
string = do
  n <- P.decimal
  _ <- P.char ':'
  String <$> P.take n

number :: Parser BValue
number = Number <$> (P.char 'i' *> P.signed P.decimal <* P.char 'e')

list :: Parser BValue
list = List <$> (P.char 'l' *> P.many' value <* P.char 'e')

dictionary :: Parser BValue
dictionary = do
  void $ P.char 'd'
  pairs <- P.many' ((,) <$> string <*> value)
  void $ P.char 'e'
  return $ Dictionary $ fixPair <$> pairs
  where fixPair (String s, v) = (s, v)

value :: Parser BValue
value = string <|> number <|> list <|> dictionary

serialize :: BValue -> ByteString
serialize (String s) = BC.pack (show $ B.length s) <> ":" <> s
serialize (Number s) = "i" <> BC.pack (show s) <> "e"
serialize (List s) = "l" <> (BC.intercalate "" . toList . fmap serialize $ s) <> "e"
serialize (Dictionary m) = "d" <> (BC.intercalate "" . Foldable.foldl folder [] $ m) <> "e"
                           where folder :: [ByteString] -> (ByteString, BValue) -> [ByteString]
                                 folder a (k, v) = (serialize (String k) <> serialize v) : a

bstring :: Traversal' BValue ByteString
bstring f (String s) = String <$> f s
bstring _ bv = pure bv

bnumber :: Traversal' BValue Integer
bnumber f (Number n) = Number <$> f n
bnumber _ bv = pure bv

blist :: Traversal' BValue BValue
blist f (List xs) = List <$> traverse f xs
blist _ bv = pure bv

bkey :: ByteString -> Traversal' BValue BValue
bkey k f bv@(Dictionary m) = case lookup k m of
                               Just v -> f v
                               Nothing -> pure bv
bkey _ _ bv = pure bv

data MetaInfo = MetaInfo {
    info     :: InfoDictionary
  , infoHash :: ByteString
  , announce :: ByteString
  }

data InfoDictionary = InfoDictionary {
    pieceLength :: Word64
  , pieces      :: ByteString
  , name        :: ByteString
  , length      :: Word64
  }

parseMetaInfo :: BValue -> Maybe MetaInfo
parseMetaInfo bv = MetaInfo
                <$> (bv ^? bkey "info" >>= parseInfoDictionary)
                <*> (hash . serialize <$> bv ^? bkey "info")
                <*> bv ^? bkey "announce" . bstring

parseInfoDictionary :: BValue -> Maybe InfoDictionary
parseInfoDictionary bv = InfoDictionary
                      <$> (fromIntegral <$> bv ^? bkey "piece length" . bnumber)
                      <*> bv ^? bkey "pieces" . bstring
                      <*> bv ^? bkey "name" . bstring
                      <*> (fromIntegral <$> bv ^? bkey "length" . bnumber)

PWP.hs

{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS_GHC -fno-warn-unused-do-bind #-}
module PWP where

import Control.Monad
import Data.Binary
import Data.Binary.Get
import Data.Binary.Put
import Data.ByteString (ByteString)
import qualified Data.ByteString as B

protocolString :: ByteString
protocolString = "BitTorrent protocol"

data PWP = KeepAlive
         | Choke
         | Unchoke
         | Interested
         | Uninterested
         | Have Word32
         | Bitfield ByteString
         | Request Word32 Word32 Word32
         | Piece Word32 Word32 ByteString
         | Cancel Word32 Word32 Word32
         deriving(Show, Eq)

instance Binary PWP where
  put KeepAlive =
    put (0 :: Word32)
  put Choke = do
    put (1 :: Word32)
    put (0 :: Word8)
  put Unchoke = do
    put (1 :: Word32)
    put (1 :: Word8)
  put Interested = do
    put (1 :: Word32)
    put (2 :: Word8)
  put Uninterested = do
    put (1 :: Word32)
    put (3 :: Word8)
  put (Have pieceId) = do
    put (5 :: Word32)
    put (4 :: Word8)
    put (pieceId :: Word32)
  put (Bitfield field) = do
    putWord32be $ fromIntegral $ 1 + B.length field
    put (5 :: Word8)
    putByteString field
  put (Request piece offset len) = do
    put (13 :: Word32)
    put (6 :: Word8)
    put piece
    put offset
    put len
  put (Piece piece offset d) = do
    put (fromIntegral $ 9 + B.length d :: Word32)
    put (7 :: Word8)
    put piece
    put offset
    putByteString d
  put (Cancel piece offset len) = do
    put (13 :: Word32)
    put (8 :: Word8)
    put piece
    put offset
    put len

  get = do
    len <- get :: Get Word32
    case len of
      0 -> return KeepAlive
      _ -> do
        messageId <- get :: Get Word8
        case messageId of
          0 -> return Choke
          1 -> return Unchoke
          2 -> return Interested
          3 -> return Uninterested
          4 -> Have <$> get
          5 -> Bitfield <$> getByteString (fromIntegral len - 1)
          6 -> Request <$> get <*> get <*> get
          7 -> Piece <$> get <*> get
                     <*> getByteString (fromIntegral len - 9)
          8 -> Cancel <$> get <*> get <*> get
          _ -> fail "incorrect!"

data BHandshake = BHandshake ByteString ByteString deriving(Show, Eq)

instance Binary BHandshake where
  put (BHandshake infoHash peerId) = do
    putWord8 (fromIntegral $ B.length protocolString)
    putByteString protocolString
    replicateM_ 8 (putWord8 0)
    putByteString infoHash
    putByteString peerId

  get = do
    protoSize <- get :: Get Word8
    replicateM_ (fromIntegral protoSize) getWord8
    skip 8
    infoHash <- getByteString 20
    peerId   <- getByteString 20
    return $ BHandshake infoHash peerId

stack.yaml

flags:
packages:
- '.'
extra-deps:
- pretty-hex-1.0
resolver: lts-3.8

torrent-blog.cabal

name:                torrent-blog
version:             0.1.0.0
homepage:            http://github.com/farnoy/torrent
author:              Jakub Okoński
maintainer:          jakub@okonski.org
cabal-version:       >=1.10

library
  exposed-modules:   Bencoding, PWP
  ghc-options:   -Wall -O -fprof-auto
  build-depends:   base >= 4.7 && < 5
                 , attoparsec
                 , binary
                 , bytestring
                 , bytestring-conversion
                 , containers
                 , cryptohash
                 , lens-family
                 , pretty-hex
  default-language:    Haskell2010