module Data.BEncode
(
BEncode(..)
, BEncodable (..), dictAssoc, Result
, (-->), (-->?), fromAssocs, fromAscAssocs
, reqKey, optKey, (>--), (>--?)
, encode, decode
, encoded, decoded
, isInteger, isString, isList, isDict
, builder, parser, decodingError, printPretty
) where
import Control.Applicative
import Control.Monad
import Data.Int
import Data.Maybe (mapMaybe)
import Data.Monoid ((<>))
import Data.Foldable (foldMap)
import Data.Traversable (traverse)
import Data.Word (Word8, Word16, Word32, Word64, Word)
import Data.Map (Map)
import qualified Data.Map as M
import Data.Set (Set)
import qualified Data.Set as S
import Data.Attoparsec.ByteString.Char8 (Parser)
import qualified Data.Attoparsec.ByteString.Char8 as P
import Data.ByteString (ByteString)
import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as BC
import qualified Data.ByteString.Lazy as Lazy
import Data.ByteString.Internal as B (c2w, w2c)
import qualified Data.ByteString.Builder as B
import qualified Data.ByteString.Builder.Prim as BP (int64Dec, primBounded)
import Data.Text (Text)
import qualified Data.Text.Encoding as T
import Data.Version
import Text.PrettyPrint hiding ((<>))
import qualified Text.ParserCombinators.ReadP as ReadP
type Dict = Map ByteString BEncode
data BEncode = BInteger !Int64
| BString !ByteString
| BList [BEncode]
| BDict Dict
deriving (Show, Read, Eq, Ord)
type Result = Either String
class BEncodable a where
toBEncode :: a -> BEncode
fromBEncode :: BEncode -> Result a
decodingError :: String -> Result a
decodingError s = Left ("fromBEncode: unable to decode " ++ s)
instance BEncodable BEncode where
toBEncode = id
fromBEncode = Right
instance BEncodable Int where
toBEncode = BInteger . fromIntegral
fromBEncode (BInteger i) = Right (fromIntegral i)
fromBEncode _ = decodingError "integer"
instance BEncodable Bool where
toBEncode = toBEncode . fromEnum
fromBEncode b = do
i <- fromBEncode b
case i :: Int of
0 -> return False
1 -> return True
_ -> decodingError "bool"
instance BEncodable Integer where
toBEncode = BInteger . fromIntegral
fromBEncode b = fromIntegral <$> (fromBEncode b :: Result Int)
instance BEncodable ByteString where
toBEncode = BString
fromBEncode (BString s) = Right s
fromBEncode _ = decodingError "string"
instance BEncodable Text where
toBEncode = toBEncode . T.encodeUtf8
fromBEncode b = T.decodeUtf8 <$> fromBEncode b
instance BEncodable a => BEncodable [a] where
toBEncode = BList . map toBEncode
fromBEncode (BList xs) = mapM fromBEncode xs
fromBEncode _ = decodingError "list"
instance BEncodable a => BEncodable (Map ByteString a) where
toBEncode = BDict . M.map toBEncode
fromBEncode (BDict d) = traverse fromBEncode d
fromBEncode _ = decodingError "dictionary"
instance (Eq a, BEncodable a) => BEncodable (Set a) where
toBEncode = BList . map toBEncode . S.toAscList
fromBEncode (BList xs) = S.fromAscList <$> traverse fromBEncode xs
fromBEncode _ = decodingError "Data.Set"
instance BEncodable () where
toBEncode () = BList []
fromBEncode (BList []) = Right ()
fromBEncode _ = decodingError "Unable to decode unit value"
instance (BEncodable a, BEncodable b) => BEncodable (a, b) where
toBEncode (a, b) = BList [toBEncode a, toBEncode b]
fromBEncode (BList [a, b]) = (,) <$> fromBEncode a <*> fromBEncode b
fromBEncode _ = decodingError "Unable to decode a pair."
instance (BEncodable a, BEncodable b, BEncodable c) => BEncodable (a, b, c) where
toBEncode (a, b, c) = BList [toBEncode a, toBEncode b, toBEncode c]
fromBEncode (BList [a, b, c]) =
(,,) <$> fromBEncode a <*> fromBEncode b <*> fromBEncode c
fromBEncode _ = decodingError "Unable to decode a triple"
instance (BEncodable a, BEncodable b, BEncodable c, BEncodable d)
=> BEncodable (a, b, c, d) where
toBEncode (a, b, c, d) = BList [ toBEncode a, toBEncode b
, toBEncode c, toBEncode d
]
fromBEncode (BList [a, b, c, d]) =
(,,,) <$> fromBEncode a <*> fromBEncode b
<*> fromBEncode c <*> fromBEncode d
fromBEncode _ = decodingError "Unable to decode a tuple4"
instance (BEncodable a, BEncodable b, BEncodable c, BEncodable d, BEncodable e)
=> BEncodable (a, b, c, d, e) where
toBEncode (a, b, c, d, e) = BList [ toBEncode a, toBEncode b
, toBEncode c, toBEncode d
, toBEncode e
]
fromBEncode (BList [a, b, c, d, e]) =
(,,,,) <$> fromBEncode a <*> fromBEncode b
<*> fromBEncode c <*> fromBEncode d <*> fromBEncode e
fromBEncode _ = decodingError "Unable to decode a tuple5"
instance BEncodable Version where
toBEncode = toBEncode . BC.pack . showVersion
fromBEncode (BString bs)
| [(v, _)] <- ReadP.readP_to_S parseVersion (BC.unpack bs)
= return v
fromBEncode _ = decodingError "Data.Version"
dictAssoc :: [(ByteString, BEncode)] -> BEncode
dictAssoc = BDict . M.fromList
data Assoc = Required ByteString BEncode
| Optional ByteString (Maybe BEncode)
(-->) :: BEncodable a => ByteString -> a -> Assoc
key --> val = Required key (toBEncode val)
(-->?) :: BEncodable a => ByteString -> Maybe a -> Assoc
key -->? mval = Optional key (toBEncode <$> mval)
mkAssocs :: [Assoc] -> [(ByteString, BEncode)]
mkAssocs = mapMaybe unpackAssoc
where
unpackAssoc (Required n v) = Just (n, v)
unpackAssoc (Optional n (Just v)) = Just (n, v)
unpackAssoc (Optional _ Nothing) = Nothing
fromAssocs :: [Assoc] -> BEncode
fromAssocs = BDict . M.fromList . mkAssocs
fromAscAssocs :: [Assoc] -> BEncode
fromAscAssocs = BDict . M.fromList . mkAssocs
reqKey :: BEncodable a => Map ByteString BEncode -> ByteString -> Result a
reqKey d key
| Just b <- M.lookup key d = fromBEncode b
| otherwise = Left ("required field `" ++ BC.unpack key ++ "' not found")
optKey :: BEncodable a => Map ByteString BEncode -> ByteString -> Result (Maybe a)
optKey d key
| Just b <- M.lookup key d
, Right r <- fromBEncode b = return (Just r)
| otherwise = return Nothing
(>--) :: BEncodable a => Map ByteString BEncode -> ByteString -> Result a
(>--) = reqKey
(>--?) :: BEncodable a => Map ByteString BEncode -> ByteString -> Result (Maybe a)
(>--?) = optKey
isInteger :: BEncode -> Bool
isInteger (BInteger _) = True
isInteger _ = False
isString :: BEncode -> Bool
isString (BString _) = True
isString _ = False
isList :: BEncode -> Bool
isList (BList _) = True
isList _ = False
isDict :: BEncode -> Bool
isDict (BList _) = True
isDict _ = False
encode :: BEncode -> Lazy.ByteString
encode = B.toLazyByteString . builder
decode :: ByteString -> Result BEncode
decode = P.parseOnly parser
decoded :: BEncodable a => ByteString -> Result a
decoded = decode >=> fromBEncode
encoded :: BEncodable a => a -> Lazy.ByteString
encoded = encode . toBEncode
builder :: BEncode -> B.Builder
builder = go
where
go (BInteger i) = B.word8 (c2w 'i') <>
BP.primBounded BP.int64Dec i <>
B.word8 (c2w 'e')
go (BString s) = buildString s
go (BList l) = B.word8 (c2w 'l') <>
foldMap go l <>
B.word8 (c2w 'e')
go (BDict d) = B.word8 (c2w 'd') <>
foldMap mkKV (M.toAscList d) <>
B.word8 (c2w 'e')
where
mkKV (k, v) = buildString k <> go v
buildString s = B.intDec (B.length s) <>
B.word8 (c2w ':') <>
B.byteString s
parser :: Parser BEncode
parser = valueP
where
valueP = do
mc <- P.peekChar
case mc of
Nothing -> fail "end of input"
Just c ->
case c of
di | di <= '9' -> BString <$> stringP
'i' -> P.anyChar *> ((BInteger <$> integerP) <* P.anyChar)
'l' -> P.anyChar *> ((BList <$> listBody) <* P.anyChar)
'd' -> do
P.anyChar
(BDict . M.fromDistinctAscList <$> many ((,) <$> stringP <*> valueP))
<* P.anyChar
t -> fail ("bencode unknown tag: " ++ [t])
listBody = do
c <- P.peekChar
case c of
Just 'e' -> return []
_ -> (:) <$> valueP <*> listBody
stringP :: Parser ByteString
stringP = do
n <- P.decimal :: Parser Int
P.char ':'
P.take n
integerP :: Parser Int64
integerP = do
c <- P.peekChar
case c of
Just '-' -> do
P.anyChar
negate <$> P.decimal
_ -> P.decimal
printPretty :: BEncode -> IO ()
printPretty = print . ppBEncode
ppBS :: ByteString -> Doc
ppBS = text . map w2c . B.unpack
ppBEncode :: BEncode -> Doc
ppBEncode (BInteger i) = int (fromIntegral i)
ppBEncode (BString s) = ppBS s
ppBEncode (BList l) = brackets $ hsep (punctuate comma (map ppBEncode l))
ppBEncode (BDict d) = braces $ vcat (punctuate comma (map ppKV (M.toAscList d)))
where
ppKV (k, v) = ppBS k <+> colon <+> ppBEncode v
instance BEncodable Word8 where
toBEncode = toBEncode . (fromIntegral :: Word8 -> Word64)
fromBEncode b = (fromIntegral :: Word64 -> Word8) <$> fromBEncode b
instance BEncodable Word16 where
toBEncode = toBEncode . (fromIntegral :: Word16 -> Word64)
fromBEncode b = (fromIntegral :: Word64 -> Word16) <$> fromBEncode b
instance BEncodable Word32 where
toBEncode = toBEncode . (fromIntegral :: Word32 -> Word64)
fromBEncode b = (fromIntegral :: Word64 -> Word32) <$> fromBEncode b
instance BEncodable Word64 where
toBEncode = toBEncode . (fromIntegral :: Word64 -> Int)
fromBEncode b = (fromIntegral :: Int -> Word64) <$> fromBEncode b
instance BEncodable Word where
toBEncode = toBEncode . (fromIntegral :: Word -> Int)
fromBEncode b = (fromIntegral :: Int -> Word) <$> fromBEncode b