{-# LANGUAGE CPP               #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RecordWildCards   #-}
module Test.TorCell(torCellTests) where

#if !MIN_VERSION_base(4,8,0)
import Control.Applicative
#endif
import Control.Monad
import Crypto.Hash
import Data.ASN1.OID
import Data.Binary.Get
import Data.Binary.Put
import Data.ByteArray(convert)
import Data.ByteString(ByteString)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as BSC
import qualified Data.ByteString.Lazy  as BSL
import Data.ByteString.Lazy(toStrict,fromStrict)
import Data.Hourglass
import Data.List
import Data.String
import Data.Word
import Data.X509
import Numeric
import Test.QuickCheck
import Test.Framework
import Test.Framework.Providers.QuickCheck2
import Test.Standard
import Tor.DataFormat.RelayCell
import Tor.DataFormat.TorAddress
import Tor.DataFormat.TorCell
import Tor.State.Credentials

instance Arbitrary TorAddress where
  arbitrary = oneof [ Hostname <$> genHostname
                    , IP4 <$> genIP4
                    , IP6 <$> genIP6
                    , return (TransientError "External transient error.")
                    , return (NontransientError "External nontransient error.")
                    ]

genHostname :: Gen String
genHostname = take 255 <$> 
                intercalate "." <$>
                  (listOf (listOf (elements ['a'..'z'])))

genIP4 :: Gen String
genIP4 = intercalate "." <$>
           (replicateM 4 (show <$> (arbitrary :: Gen Word8)))

genIP6 :: Gen String
genIP6 = do x <- genIP6'
            return ("[" ++ intercalate ":" x ++ "]")
 where
  genIP6'     = map showHex' <$> 
                        replicateM 8 (arbitrary :: Gen Word16)

prop_TorAddrSerial :: TorAddress -> Bool
prop_TorAddrSerial  = serialProp getTorAddress putTorAddress

data TorAddressBS = TABS ByteString TorAddress
 deriving (Show, Eq)

instance Arbitrary TorAddressBS where
  arbitrary = oneof [ do x <- replicateM 4 arbitrary
                         let str = intercalate "." (map show x)
                             bstr = BS.pack x
                         return (TABS bstr (IP4 str))
                    , do x <- replicateM 16 arbitrary
                         let bstr = BSL.pack x
                             xs   = runGet (replicateM 8 getWord16be) bstr
                             str  = "[" ++ intercalate ":" (map showHex' xs) ++ "]"
                         return (TABS (toStrict bstr) (IP6 str))
                    ]

prop_TorAddrBSSerial :: TorAddressBS -> Bool
prop_TorAddrBSSerial (TABS bstr x) = bstr == torAddressByteString x

showHex' :: (Show a, Integral a) => a -> String
showHex' x = showHex x ""

instance Arbitrary ExtendSpec where
  arbitrary = oneof [ ExtendIP4 <$> genIP4 <*> arbitrary
                    , ExtendIP6 <$> genIP6 <*> arbitrary
                    , ExtendDigest <$>
                        (BSC.pack <$>
                           replicateM 20 (elements "abcdef0123456789"))
                    ]

prop_ExtendSpecSerial :: ExtendSpec -> Bool
prop_ExtendSpecSerial = serialProp getExtendSpec putExtendSpec

instance Arbitrary DestroyReason where
  arbitrary = elements [NoReason, TorProtocolViolation, InternalError,
                        RequestedDestroy, NodeHibernating, HitResourceLimit,
                        ConnectionFailed, ORIdentityIssue, ORConnectionClosed,
                        Finished, CircuitConstructionTimeout, CircuitDestroyed,
                        NoSuchService]

prop_DestroyReasonSerial1 :: DestroyReason -> Bool
prop_DestroyReasonSerial1 = serialProp getDestroyReason putDestroyReason

prop_DestroyReasonSerial2 :: Word8 -> Bool
prop_DestroyReasonSerial2 x =
  [x] == BSL.unpack (runPut (putDestroyReason
                       (runGet getDestroyReason (BSL.pack [x]))))

instance Arbitrary RelayEndReason where
  arbitrary = oneof [ ReasonExitPolicy <$> (IP4 <$> genIP4) <*> arbitrary
                    , ReasonExitPolicy <$> (IP6 <$> genIP6) <*> arbitrary
                    , elements [ReasonMisc, ReasonResolveFailed,
                        ReasonConnectionRefused, ReasonDestroyed, ReasonDone,
                        ReasonTimeout, ReasonNoRoute, ReasonHibernating,
                        ReasonInternal, ReasonResourceLimit,
                        ReasonConnectionReset, ReasonTorProtocol,
                        ReasonNotDirectory ]
                    ]

prop_RelayEndRsnSerial :: RelayEndReason -> Bool
prop_RelayEndRsnSerial rsn =
  let bstr = runPut (putRelayEndReason rsn)
      len  = case rsn of
               ReasonExitPolicy (IP4 _) _ -> 9
               ReasonExitPolicy (IP6 _) _ -> 21
               _                          -> 1
      rsn' = runGet (getRelayEndReason len) bstr
  in rsn == rsn'

instance Arbitrary RelayCell where
  arbitrary =
   oneof [ RelayBegin <$> arbitrary <*> legalTorAddress True
                      <*> arbitrary <*> arbitrary <*> arbitrary
                      <*> arbitrary
         , RelayData <$> arbitrary
                     <*> ((BS.pack . take 503) <$> arbitrary)
         , RelayEnd <$> arbitrary <*> arbitrary
         , RelayConnected <$> arbitrary <*> legalTorAddress False
                          <*> arbitrary
         , RelaySendMe <$> arbitrary
         , RelayExtend <$> arbitrary <*> genIP4
                       <*> arbitrary <*> arbitraryBS 186 <*> arbitraryBS 20
         , RelayExtended <$> arbitrary
                         <*> (BS.pack <$> replicateM 148 arbitrary)
         , RelayTruncate <$> arbitrary
         , RelayTruncated <$> arbitrary <*> arbitrary
         , RelayDrop <$> arbitrary
         , RelayResolve <$> arbitrary
                        <*> (filter (/= '\0') <$> arbitrary)
         , do strm <- arbitrary
              vals <- listOf $ do x <- legalTorAddress True
                                  y <- arbitrary
                                  return (x,y)
              return (RelayResolved strm vals)
         , RelayBeginDir <$> arbitrary
         , RelayExtend2 <$> arbitrary <*> arbitrary <*> arbitrary
                        <*> (BS.pack <$> arbitrary)
         , RelayExtended2 <$> arbitrary <*> (BS.pack <$> arbitrary)
         , RelayEstablishIntro <$> arbitrary <*> arbitraryBS 128
                               <*> arbitraryBS 20 <*> arbitraryBS 128
         , RelayEstablishRendezvous <$> arbitrary <*> arbitraryBS 20
         , RelayIntroduce1 <$> arbitrary <*> arbitraryBS 20
                           <*> (BS.pack <$> arbitrary)
         , RelayIntroduce2 <$> arbitrary <*> (BS.pack <$> arbitrary)
         , RelayRendezvous1 <$> arbitrary <*> arbitraryBS 20
                            <*> arbitraryBS 128 <*> arbitraryBS 20
         , RelayRendezvous2 <$> arbitrary <*> arbitraryBS 128
                            <*> arbitraryBS 20
         , RelayIntroEstablished <$> arbitrary
         , RelayRendezvousEstablished <$> arbitrary
         , RelayIntroduceAck <$> arbitrary
         ]

legalTorAddress :: Bool -> Gen TorAddress
legalTorAddress allowHostname =
  do x <- arbitrary
     case x of
       Hostname ""                 -> legalTorAddress allowHostname
       Hostname _  | allowHostname -> return x
       IP4      "0.0.0.0"          -> legalTorAddress allowHostname
       IP4      _                  -> return x
       IP6      _                  -> return x
       _                           -> legalTorAddress allowHostname

prop_RelayCellSerial :: RelayCell -> Property
prop_RelayCellSerial x =
  let (_, gutsBS) = runPutM (putRelayCellGuts x)
      bstr        = runPut (putRelayCell (BS.replicate 4 0) x)
      (_, y)      = runGet getRelayCell bstr
  in (BSL.length gutsBS <= (509 - 11)) ==> (x == y)

prop_RelayCellDigestWorks1 :: RelayCell -> Property
prop_RelayCellDigestWorks1 x =
  let (_, gutsBS) = runPutM (putRelayCellGuts x)
      (bstr, _)   = renderRelayCell hashInit x
      (x',   _)   = runGet (parseRelayCell hashInit) (fromStrict bstr)
  in (BSL.length gutsBS <= (509 - 11)) ==> (x == x')

prop_RelayCellDigestWorks2 :: NonEmptyList RelayCell -> Property
prop_RelayCellDigestWorks2 xs =
  let mxSize = maximum (map putGuts (getNonEmpty xs))
      xs'    = runCheck hashInit hashInit (getNonEmpty xs)
  in (mxSize <= (509 - 11)) ==> (getNonEmpty xs == xs')
 where
  putGuts x =
    let (_, gutsBS) = runPutM (putRelayCellGuts x)
    in BSL.length gutsBS
  runCheck _ _ [] = []
  runCheck rstate pstate (f:rest) =
    let (bstr, rstate') = renderRelayCell rstate f
        (f',   pstate') = runGet (parseRelayCell pstate) (fromStrict bstr)
    in f' : runCheck rstate' pstate' rest

instance Arbitrary HandshakeType where
  arbitrary = elements [TAP, Reserved, NTor]

prop_HandTypeSerial1 :: HandshakeType -> Bool
prop_HandTypeSerial1 = serialProp getHandshakeType putHandshakeType

prop_HandTypeSerial2 :: Word16 -> Bool
prop_HandTypeSerial2 x =
  let ht = runGet getHandshakeType (runPut (putWord16be x))
  in runPut (putWord16be x) == runPut (putHandshakeType ht)

instance Arbitrary (SignedExact Certificate) where
  arbitrary =
    do certVersion      <- arbitrary
       certSerial       <- arbitrary
       certIssuerDN     <- arbitrary
       certSubjectDN    <- arbitrary
       hashAlg          <- elements [HashSHA1, HashSHA256, HashSHA384]

       g  <- arbitraryRNG
       let (pub, _, _) = generateKeyPair g 1024

       let keyAlg           = PubKeyALG_RSA -- FIXME?
           certSignatureAlg = SignatureALG hashAlg keyAlg
           certValidity     = (timeFromElapsed (Elapsed (Seconds 257896558)),
                               timeFromElapsed (Elapsed (Seconds 2466971758)))
           certPubKey       = PubKeyRSA pub
           certExtensions   = Extensions Nothing

       let baseCert = Certificate{ .. }
           sigfun   = case hashAlg of
                        HashSHA1   -> wrapSignatureAlg certSignatureAlg sha1
                        HashSHA224 -> wrapSignatureAlg certSignatureAlg sha224
                        HashSHA256 -> wrapSignatureAlg certSignatureAlg sha256
                        HashSHA384 -> wrapSignatureAlg certSignatureAlg sha384
                        HashSHA512 -> wrapSignatureAlg certSignatureAlg sha512
                        _          -> error "INTERNAL WEIRDNESS"
       let (signedCert, _) = objectToSignedExact sigfun baseCert
       return signedCert

newtype ReadableStr = ReadableStr { unReadableStr :: String }

instance Show ReadableStr where
  show = show . unReadableStr

instance Arbitrary ReadableStr where
  arbitrary =
    do len <- choose (1, 256)
       str <- replicateM len (elements printableChars)
       return (ReadableStr str)
   where printableChars = ['a'..'z'] ++ ['A'..'Z'] ++ ['_','.',' ']

instance Arbitrary DistinguishedName where
  arbitrary =
    do cn <- unReadableStr <$> arbitrary
       co <- unReadableStr <$> arbitrary
       og <- unReadableStr <$> arbitrary
       ou <- unReadableStr <$> arbitrary
       return (DistinguishedName [
                 (getObjectID DnCommonName,       fromString cn)
               , (getObjectID DnCountry,          fromString co)
               , (getObjectID DnOrganization,     fromString og)
               , (getObjectID DnOrganizationUnit, fromString ou)
               ])

wrapSignatureAlg :: SignatureALG ->
                    (ByteString -> ByteString) ->
                    ByteString ->
                    (ByteString, SignatureALG, ())
wrapSignatureAlg name sha bstr =
  let hashed   = convert (sha bstr)
  in (hashed, name, ())

sha1 :: ByteString -> ByteString
sha1 = convert . hashWith SHA1

sha224 :: ByteString -> ByteString
sha224 = convert . hashWith SHA224

sha256 :: ByteString -> ByteString
sha256 = convert . hashWith SHA256

sha384 :: ByteString -> ByteString
sha384 = convert . hashWith SHA384

sha512 :: ByteString -> ByteString
sha512 = convert . hashWith SHA512


instance Arbitrary TorCert where
  arbitrary = oneof [ LinkKeyCert     <$> arbitrary
                    , RSA1024Identity <$> arbitrary
                    , RSA1024Auth     <$> arbitrary
                    ]

prop_torCertSerial :: TorCert -> Bool
prop_torCertSerial = serialProp getTorCert putTorCert

torCellTests :: Test
torCellTests =
  testGroup "TorCell Serialization" [
    testProperty "TorAddress round-trips" prop_TorAddrSerial
  , testProperty "TorAddress makes sensible ByteStrings" prop_TorAddrBSSerial
  , testProperty "ExtendSpec serializes" prop_ExtendSpecSerial
  , testProperty "DestroyReason serializes (check #1)" prop_DestroyReasonSerial1
  , testProperty "DestroyReason serializes (check #2)" prop_DestroyReasonSerial2
  , testProperty "HandshakeType serializes (check #1)" prop_HandTypeSerial1
  , testProperty "HandshakeType serializes (check #2)" prop_HandTypeSerial2
  , testProperty "RelayEndReason serializes" prop_RelayEndRsnSerial
  , testProperty "RelayCell serializes" prop_RelayCellSerial
  , testProperty "RelayCell serializes w/ digest" prop_RelayCellDigestWorks1
  , testProperty "RelayCell serializes w/ digest" prop_RelayCellDigestWorks2
  , testProperty "Tor certificates serialize" prop_torCertSerial
  ]

