Torrent.hs

{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings          #-}
{-# LANGUAGE TemplateHaskell            #-}
-- |

module Data.Torrent where

import           Control.Applicative
import           Control.Lens
import           Data.AttoBencode
import           Data.ByteString (ByteString)
import qualified Data.ByteString as BS hiding (unpack)
import qualified Data.ByteString.Char8 as BS (unpack)
import           Data.Maybe (fromMaybe)
import           Data.Serialize hiding (decode)
import           Data.Text (Text)
import qualified Data.Text as Text
import           Data.Text.Encoding (decodeUtf8)
import           Network.URI

-- | Get an optional key
d .:? k = pure (d .: k)

-- | Provide a default for a required key
x .|| y = pure (fromMaybe y x)

instance FromBencode URI where
  fromBencode (BString s) = parseURI (BS.unpack s)
  fromBencode _ = Nothing

instance FromBencode Bool where
  fromBencode (BInt 0) = Just False
  fromBencode (BInt 1) = Just True
  fromBencode _ = Nothing

instance FromBencode Text where
  fromBencode (BString s) = Just (decodeUtf8 s)
  fromBencode _ = Nothing

newtype InfoHash = InfoHash { _getInfoHash :: ByteString }
  deriving (Eq, Ord, Show, FromBencode)

instance Serialize InfoHash where
  put = putByteString . _getInfoHash
  get = InfoHash <$> getBytes 20

newtype Pieces = Pieces [InfoHash]

instance Show Pieces where
  show (Pieces xs) = show (length xs) ++ " Pieces"

instance FromBencode Pieces where
  fromBencode (BString s) = Just (Pieces (InfoHash <$> chunksOf 20 s))
    where chunksOf n s | BS.null s = []
                       | otherwise = BS.take n s : chunksOf n (BS.drop n s)
  fromBencode _ = Nothing

data FileInfo = FileInfo
  { _fPath :: [Text]
  , _fLength :: Integer
  } deriving (Show)

makeLenses ''FileInfo

instance FromBencode FileInfo where
  fromBencode (BDict d) = FileInfo <$> d .: "path" <*> d .: "length"
  fromBencode _ = Nothing

data Info = Info
  { _iName       :: !Text
  , _pieces      :: !Pieces
  , _pieceLength :: !Integer
  , _length      :: !(Maybe Integer)
  , _files       :: !(Maybe [FileInfo])
  } deriving (Show)

makeLenses ''Info

instance FromBencode Info where
  fromBencode (BDict d) = Info <$> d .: "name"
                               <*> d .: "pieces"
                               <*> d .: "piece length"
                               <*> d .:? "length"
                               <*> d .:? "files"

  fromBencode _ = Nothing

data Torrent = Torrent
  { _announce     :: !(Maybe URI)
  , _announceList :: ![[URI]]
  , _info         :: !Info
  , _comment      :: !(Maybe Text)
  , _private      :: !Bool
  , _path         :: ![Text]
  } deriving (Show)

instance FromBencode Torrent where
  fromBencode (BDict d) = Torrent <$> d .:? "announce"
                                  <*> d .: "announce-list" .|| []
                                  <*> d .: "info"
                                  <*> d .:? "comment"
                                  <*> d .: "private"       .|| False
                                  <*> d .: "path"          .|| []
  fromBencode _ = Nothing

makeLenses ''Torrent

class HasName a where
  name :: Lens' a Text

instance HasName Info where
  name = iName

instance HasName Torrent where
  name = info . iName

decodeTorrent :: BS.ByteString -> Maybe Torrent
decodeTorrent = decode