{-# LANGUAGE OverloadedStrings #-}

module Main where

import           Helper ((-), log)
import           Prelude hiding ((-), log)
import qualified Prelude as Prelude

import           Control.Concurrent.Async (async, waitAnyCatchCancel, waitEitherCancel)
import           Control.Lens
import           Control.Monad (replicateM_, join, when)
import           Control.Monad.IO.Class (MonadIO(..))
import           Control.Monad.Trans.State.Strict (runStateT)
import           Data.Attoparsec.ByteString (parse, maybeResult)
import           Data.ByteString (ByteString)
import qualified Data.ByteString as B
import qualified Data.ByteString.Builder as Builder
import qualified Data.ByteString.Char8 as BC
import           Data.Monoid ((<>))
import qualified Data.Text as T
import           Data.Text.Lens
import           Data.Text.Strict.Lens (utf8)
import           Network.Simple.TCP (send)
import qualified Network.Socket as NS
import           Pipes (runEffect, (>->), Pipe, for, yield, cat, Producer, Consumer)
import qualified Pipes.Attoparsec as PA
import           Pipes.Network.TCP.Safe
import qualified Pipes.Prelude as P
import           Pipes.Safe
import           System.Random (randomIO)
import           Options.Generic
import           System.IO (hSetBuffering, stdin, stdout, BufferMode(LineBuffering))

import           NetworkHelper (setSocketCommon)
import           Options (toConfig)
import           Parser
import           Type


obfsEncode :: (MonadIO m) => Bool -> Integer -> Pipe ByteString ByteString m a
obfsEncode debug randomness = do
  for cat - \x -> do
    let payloadLen = (fromIntegral - B.length x)
    if (payloadLen >= fromIntegral randomness)
      then do
        let paddingLen = 0
        let chunk =
              view strict - Builder.toLazyByteString -
                Builder.word32BE paddingLen
                <> Builder.byteString
                      (BC.replicate (fromIntegral paddingLen) - '.')
                <> Builder.word32BE payloadLen
                <> Builder.byteString x

        when debug - log - "E: " <> (T.pack - show - B.length chunk)
        yield chunk
      else
        randomLenCheckLoop x
  where
    randomLenCheckLoop payload = do
      let payloadLen = (fromIntegral - B.length payload)
      randomLen <- liftIO - randomIO <&> (`mod` fromIntegral randomness)
      if (payloadLen <= randomLen)
        then
          do
            let paddingLen = randomLen Prelude.- payloadLen
            let chunk =
                  view strict - Builder.toLazyByteString -
                    Builder.word32BE paddingLen
                    <> Builder.byteString
                        (BC.replicate (fromIntegral paddingLen) - '.')
                    <> Builder.word32BE payloadLen
                    <> Builder.byteString payload

            when debug - log - "E: " <> (T.pack - show - B.length chunk)
            yield chunk
        else
          do
            let paddingLen = 0
                (now, next) = B.splitAt (fromIntegral randomLen) payload

            let chunk =
                  view strict - Builder.toLazyByteString -
                    Builder.word32BE paddingLen
                    <> Builder.byteString
                        (BC.replicate (fromIntegral paddingLen) - '.')
                    <> Builder.word32BE (fromIntegral - B.length now)
                    <> Builder.byteString now

            when debug - log - "E: " <> (T.pack - show - B.length chunk)
            yield chunk

            randomLenCheckLoop next


obfsDecode :: (MonadIO m) => Producer ByteString m a -> (ByteString -> m ()) -> m ()
obfsDecode pull sink = loop pull
  where
    loop x = do
      (r, next) <- runStateT (PA.parse obfsParser) x
      case r of
        Just (Right result@(_,_,_,payload)) -> do
          -- log - view packed - show result
          sink payload
          loop next
        _ -> pure ()

main :: IO ()
main = do
  hSetBuffering stdin LineBuffering
  hSetBuffering stdout LineBuffering

  config <- getRecord "a TCP tunnel with packet length obfuscation" <&> toConfig

  log - T.pack - show config

  let _mtu = fromIntegral - config ^. mtu + (-8) -- 2 word32be
      _timeout = fromIntegral - config ^. timeoutInSeconds * 1000000

  let _localHost = config ^. localHost . re packed
      _localPort = show - config ^. localPort

      _remoteHost = config ^. remoteHost . re packed
      _remotePort = show - config ^. remotePort

      _fromSocket = fromSocketTimeout _timeout
      _toSocket = toSocketTimeout _timeout

      _debug = config ^. debug

  localThread <- pure -
    runSafeT . runEffect - serve (Host _localHost) _localPort - \(localSock, localSockAddr) -> do
      log - "local accepted: " <> view packed (show localSockAddr)
      setSocketCommon localSock

      runSafeT . runEffect -
        connect _remoteHost _remotePort - \(remoteSock, remoteSockAddr) -> do
          liftIO - setSocketCommon remoteSock

          let
              localPull = _fromSocket localSock _mtu

              remotePull = _fromSocket remoteSock _mtu
              remotePush = _toSocket remoteSock


          sendThread <- liftIO - async - runEffect - do
            localPull >-> obfsEncode _debug (config ^. randomnessInBytes)
                      >-> remotePush
            liftIO - NS.shutdown remoteSock NS.ShutdownSend

          recvThread <- liftIO - async - do
            obfsDecode remotePull - send localSock
            liftIO - NS.shutdown localSock NS.ShutdownSend

          liftIO - waitEitherCancel sendThread recvThread

          pure ()

  remoteThread <- pure -
    runSafeT . runEffect - serve (Host _remoteHost) _remotePort - \(remoteSock, remoteSockAddr) -> do
      log - "remote accepted: " <> view packed (show remoteSockAddr)
      setSocketCommon remoteSock

      let _forwardHost = config ^. forwardHost . re packed
          _forwardPort = show - config ^. forwardPort

      runSafeT . runEffect -
        connect _forwardHost _forwardPort - \(forwardSock, forwardSockAddr) -> do
          liftIO - setSocketCommon forwardSock

          let
              remotePull = _fromSocket remoteSock _mtu
              remotePush = _toSocket remoteSock

              forwardPull = _fromSocket forwardSock _mtu

          sendThread <- liftIO - async - do
            obfsDecode remotePull - send forwardSock
            liftIO - NS.shutdown forwardSock NS.ShutdownSend

          recvThread <- liftIO - async - runEffect - do
            forwardPull >-> obfsEncode _debug (config ^. randomnessInBytes)
                        >-> remotePush
            liftIO - NS.shutdown remoteSock NS.ShutdownSend

          liftIO - waitEitherCancel sendThread recvThread

          pure ()


  case config ^. role of
    Local -> join localThread
    Remote -> join remoteThread
    Both -> do
      local <- async - join localThread
      remote <- async - join remoteThread
      waitEitherCancel local remote
      pure ()

  pure ()