BuilderAndParser.hs

{-# LANGUAGE OverloadedStrings #-}

module Network.MoeSocks.BuilderAndParser where

import Control.Lens
import Data.Attoparsec.ByteString
import Data.Binary
import Data.Binary.Put
import Data.Monoid
import Data.Maybe
import Data.Text.Lens
import Data.Text.Strict.Lens (utf8)
import Safe (readMay)
import Network.MoeSocks.Constant
import Network.MoeSocks.Helper
import Network.MoeSocks.Type
import Prelude hiding ((-), take)
import qualified Prelude as P
import qualified Data.ByteString as S
import qualified Data.ByteString.Builder as B
import Network.Socket

socksVersion :: Word8
socksVersion = 5

socksHeader :: Parser Word8
socksHeader = word8 socksVersion

greetingParser :: Parser ClientGreeting
greetingParser = do
  socksHeader
  let maxNoOfMethods = 5
  _numberOfAuthenticationMethods <- satisfy (<= maxNoOfMethods)

  ClientGreeting <$>
    count (fromIntegral _numberOfAuthenticationMethods) anyWord8

greetingReplyBuilder :: B.Builder
greetingReplyBuilder =  B.word8 socksVersion
                     <> B.word8 _No_authentication

portParser :: Parser Int
portParser = do
  __portNumberPair <- (,) <$> anyWord8 <*> anyWord8
  pure - portPairToInt __portNumberPair

portBuilder :: (Integral i) => i -> B.Builder
portBuilder i =
  let _i = fromIntegral i :: Word16 
  in
  foldMapOf both B.word8 -
    (decode - runPut - put _i :: (Word8, Word8))

requestParser :: Parser ClientRequest
requestParser = do
  __connectionType <- choice
      [
        TCP_IP_stream_connection <$ word8 1 
      , TCP_IP_port_binding <$ word8 2
      , UDP_port <$ word8 3
      ]

  word8 _ReservedByte
  __addressType <- addressTypeParser
  __portNumber <- portParser
  pure - 
          ClientRequest
            __connectionType
            __addressType 
            __portNumber

connectionParser :: Parser ClientRequest 
connectionParser = do
  socksHeader
  requestParser

sockAddr_To_Pair :: SockAddr -> (AddressType, Port)
sockAddr_To_Pair (SockAddrInet _port _host) =
                        let 
                            _r@(_a, _b, _c, _d) = decode . 
                                    runPut - put - _host
                                      :: (Word8, Word8, Word8, Word8)
                        in
                        ( IPv4_address - flip4 _r
                        , fromIntegral _port
                        )
sockAddr_To_Pair (SockAddrInet6 _port _ _host _) =
                        let 
                            _r@(_a, _b, _c, _d) = decode . 
                                    runPut - put - _host
                                      :: (Word32, Word32, Word32, Word32)
                        in
                        ( IPv6_address - flip4 _r
                        , fromIntegral _port
                        )
sockAddr_To_Pair (SockAddrUnix x) = 
                        let
                            _host = P.takeWhile (/= ':') x :: String 
                            _port = x & reverse & P.takeWhile (/= ':') 
                                      & reverse
                        in
                        ( Domain_name - (_host & review _Text)
                        , fromMaybe 0 - readMay _port
                        )
sockAddr_To_Pair sa = error - "SockAddrCan not implemented: " <> show sa 


connectionReplyBuilder :: SockAddr -> B.Builder
connectionReplyBuilder aSockAddr = 
  let _r@(__addressType, _port) = sockAddr_To_Pair aSockAddr
  in


  let _b =
              B.word8 socksVersion
          <>  B.word8 _Request_Granted 
          <>  B.word8 _ReservedByte
          <>  addressTypeBuilder __addressType
          <>  portBuilder _port
  in
  _b

addressTypeBuilder :: AddressType -> B.Builder
addressTypeBuilder aAddressType = 
  case aAddressType of
    IPv4_address _address -> 
                          B.word8 1
                       <> foldMapOf each B.word8 _address
    Domain_name x ->   
                          B.word8 3
                       <> B.word8 (fromIntegral (S.length (review utf8 x)))
                       <> B.byteString (review utf8 x)

    IPv6_address _address ->  
                          B.word8 4
                       <> foldMapOf each B.word32BE _address



connectionType_To_Word8 :: ConnectionType -> Word8
connectionType_To_Word8 TCP_IP_stream_connection = 1
connectionType_To_Word8 TCP_IP_port_binding = 2
connectionType_To_Word8 UDP_port = 3



requestBuilder :: ClientRequest -> B.Builder
requestBuilder aClientRequest = 
      B.word8 (connectionType_To_Word8 - aClientRequest ^. connectionType)
  <>  B.word8 _ReservedByte
  <>  addressTypeBuilder (aClientRequest ^. addressType)
  <>  portBuilder (aClientRequest ^. portNumber)

shadowsocksRequestBuilder :: ClientRequest -> B.Builder
shadowsocksRequestBuilder aClientRequest =
      addressTypeBuilder (aClientRequest ^. addressType)
  <>  portBuilder (aClientRequest ^. portNumber)

anyWord32be :: Parser Word32
anyWord32be = do
  _b <- count 4 anyWord8
  pure - decode - runPut - put _b

addressTypeParser :: Parser AddressType
addressTypeParser = choice
  [
    IPv4_address <$>  do
                        word8 1
                        _a <- anyWord8
                        _b <- anyWord8
                        _c <- anyWord8
                        _d <- anyWord8 
                        pure - (_a, _b, _c, _d)
  
  , Domain_name <$>   do 
                        word8 3
                        let maxDomainNameLength = 32
                        _nameLength <- satisfy (<= maxDomainNameLength)
                        view utf8 <$> (take - fromIntegral _nameLength)

  , IPv6_address <$>  do
                        word8 4 
                        _a <- anyWord32be
                        _b <- anyWord32be
                        _c <- anyWord32be
                        _d <- anyWord32be 
                        pure - (_a, _b, _c, _d)
  ]


shadowsocksRequestParser :: Parser ClientRequest
shadowsocksRequestParser = do
  __addressType <- addressTypeParser
  __portNumber <- portParser

  pure - 
          ClientRequest
            TCP_IP_stream_connection
            __addressType 
            __portNumber