{-# LANGUAGE ViewPatterns #-} module Main ( main ) where import Pipes import Pipes.Network.TCP import qualified Pipes.Binary as PB import Control.Monad import qualified Control.Monad.Trans.State.Strict as S import Control.Applicative import Crypto.Random import Crypto.PubKey.RSA import Crypto.PubKey.RSA.PKCS15 import Crypto.Cipher.Types import Crypto.Cipher.AES import qualified Data.ByteString as BS import Data.Binary import Data.Binary.Get import Data.Binary.Put import Data.Binary.IEEE754 import qualified Data.Text as T import Data.Text.Encoding import Data.X509 import Data.ASN1.Types import Data.ASN1.BinaryEncoding import Data.ASN1.Encoding import qualified System.IO.Error as E getTextL :: Get T.Text getTextL = do len <- getWord16be bytes <- getByteString . fromIntegral $ len * 2 return $ decodeUtf16BE bytes putTextL :: T.Text -> Put putTextL p = do let bs = encodeUtf16BE p putWord16be $ fromIntegral $ BS.length bs `quot` 2 putByteString bs getBSL :: Get BS.ByteString getBSL = do len <- getWord16be bytes <- getByteString $ fromIntegral len return bytes putBSL :: BS.ByteString -> Put putBSL p = do putWord16be $ fromIntegral $ BS.length p putByteString p -- TODO: Remove the Binary instances for these data types which are one-way -- decoding/encoding respectively. data InPacket = InPKeepAlive Word32 | InPHandshake Word8 T.Text T.Text Word32 | InPEncryptionResponse BS.ByteString BS.ByteString | InPClientStatus Word8 deriving (Eq, Show) instance Binary InPacket where get = do opcode <- getWord8 case opcode of 0x00 -> InPKeepAlive <$> getWord32be 0x02 -> InPHandshake <$> get <*> getTextL <*> getTextL <*> getWord32be 0xCD -> InPClientStatus <$> getWord8 0xFC -> InPEncryptionResponse <$> getBSL <*> getBSL _ -> error $ "Unsupported packet with opcode " ++ show opcode put = error "Unsupported" data OutPacket = OutPKeepAlive Word32 | OutPEncryptionRequest T.Text BS.ByteString BS.ByteString | OutPEncryptionResponse BS.ByteString BS.ByteString | OutPLoginRequest Word32 T.Text Word8 Word8 Word8 Word8 Word8 | OutPSpawnPosition Word32 Word32 Word32 | OutPPositionAndLook Double Double Double Double Float Float Word8 deriving (Eq, Show) instance Binary OutPacket where get = error "Unsupported" put p = case p of OutPKeepAlive ka -> do putWord8 0x00 putWord32be ka OutPEncryptionRequest sid cert vtok -> do putWord8 0xFD putTextL sid putBSL cert putBSL vtok OutPEncryptionResponse ss vt -> do putWord8 0xFC putBSL ss putBSL vt OutPLoginRequest eid lev gm dim diff resv maxp -> do putWord8 0x01 putWord32be eid putTextL lev putWord8 gm putWord8 dim putWord8 diff putWord8 resv putWord8 maxp OutPSpawnPosition x y z -> do putWord8 0x06 putWord32be x putWord32be y putWord32be z OutPPositionAndLook x ys sy z yaw pitch og -> do putWord8 0x0D putFloat64be x putFloat64be ys putFloat64be sy putFloat64be z putFloat32be yaw putFloat32be pitch putWord8 og -- TODO: Use >> in pipes chain to make it very hard to break processPacket :: MonadIO m => Pipe (t, InPacket) OutPacket (S.StateT SessionState m) a processPacket = forever $ do (_, p) <- await liftIO . putStrLn $ "Recieved a " ++ show p ++ " packet" (SessionState globs locs) <- lift S.get case p of InPKeepAlive{} -> error "Got a KeepAlive packet :(" InPHandshake{} -> do let (verif, _) = cprgGenerateWithEntropy 4 (gsRNG globs) -- Lens needed here _ <- lift $ S.modify (\o@(ssLocalState -> s) -> o { ssLocalState = s { lsVerifyBytes = verif } }) yield $ OutPEncryptionRequest (gsServerName globs) (gsASN1Cert globs) verif InPEncryptionResponse css vt -> case decrypt Nothing (gsPrivateKey globs) vt of Left _ -> error "Verification bytes decryption error" Right b -> do when (lsVerifyBytes locs /= b) $ error "Verification bytes do not match" let sec = either (error "S-Secret decryption error") id $ decrypt Nothing (gsPrivateKey globs) css aes = initAES sec eiv = maybe (error "Bad AES IV") id $ makeIV sec yield $ OutPEncryptionResponse BS.empty BS.empty _ <- lift $ S.modify (\o@(ssLocalState -> s) -> o { ssLocalState = s { lsCrypto = Just aes , lsCryptoIV = eiv , lsVerifyBytes = undefined } }) return () InPClientStatus{} -> do --yield $ OutPLoginRequest 1298 (T.pack "default") 0 0 1 0 8 yield $ OutPKeepAlive 1337 --yield $ OutPSpawnPosition 0 0 0 --yield $ OutPPositionAndLook 0 0 0 0 0 0 1 protocol :: MonadIO m => Socket -> Int -> Effect (S.StateT SessionState m) () protocol sock tout = let helper f bs = do (crypto, iv) <- lift $ S.gets $ \x -> let s = ssLocalState x in (lsCrypto s, lsCryptoIV s) case crypto of Just aes -> yield $ f aes iv bs Nothing -> yield bs recvS = for (fromSocketTimeout tout sock 4096) $ helper cfb8Decrypt sendS = for cat (helper cfb8Encrypt) >-> toSocketTimeout tout sock in void $ PB.decodeMany recvS >-> processPacket >-> for cat PB.encode >-> sendS data SessionState = SessionState { ssGlobalState :: GlobalState , ssLocalState :: LocalState } data GlobalState = GlobalState { gsServerName :: T.Text , gsEntropyPool :: EntropyPool , gsASN1Cert :: BS.ByteString , gsPublicKey :: PublicKey , gsPrivateKey :: PrivateKey , gsRNG :: SystemRNG } data LocalState = LocalState { lsCrypto :: Maybe AES , lsCryptoIV :: IV AES , lsVerifyBytes :: BS.ByteString } defaultLocalState :: LocalState defaultLocalState = LocalState { lsCrypto = Nothing , lsCryptoIV = undefined , lsVerifyBytes = undefined } main :: IO () main = withSocketsDo $ do entrop <- createEntropyPool let rng = cprgCreate entrop (pubKey, privKey) = fst $ generate rng 128 0x10001 serverState = GlobalState { gsServerName = T.pack "Nw]1QaC do putStrLn $ "Connection established from " ++ show addr let sockTimeout = 30 * 10 ^ (6 :: Int) E.catchIOError (void $ S.runStateT (runEffect $ protocol sock sockTimeout) $ SessionState { ssGlobalState = serverState , ssLocalState = defaultLocalState }) $ \_ -> do putStrLn $ "Client at " ++ show addr ++ " disconnected."