{-# LANGUAGE OverloadedStrings #-}

module Main where

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

import           Control.Concurrent.Async (async, waitAnyCatchCancel, waitEitherCancel)
import           Control.Lens
import           Control.Monad (replicateM_, join)
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           Options (toConfig)
import           Parser
import           Type


obfsEncode :: (MonadIO m) => Integer -> Integer -> Pipe ByteString ByteString m a
obfsEncode bound randomness = do
  for cat - \x -> do
    let payloadLen = (fromIntegral - B.length x)

    randomLen <-
      if (payloadLen > fromIntegral bound)
        then pure 0
        else liftIO - randomIO <&> (`mod` fromIntegral randomness)

    yield - view strict - Builder.toLazyByteString - Builder.word32BE randomLen
    replicateM_ (fromIntegral randomLen) - yield - BC.singleton 'a'

    yield - view strict - Builder.toLazyByteString - Builder.word32BE payloadLen
    yield x


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

main :: IO ()
main = do
  config <- getRecord "a TCP tunnel with packet length obfuscation" <&> toConfig

  log - T.pack - show config

  let mtu = 4096
      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

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

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

          let
              localPull = _fromSocket localSock mtu

              remotePull = _fromSocket remoteSock mtu
              remotePush = _toSocket remoteSock


          sendThread <- liftIO - async - runEffect - do
            localPull >-> obfsEncode (config ^. bound) (config ^. randomness)
                      >-> 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)

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

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

          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 (config ^. bound) (config ^. randomness)
                        >-> remotePush
            liftIO - NS.shutdown remoteSock NS.ShutdownSend

          liftIO - waitEitherCancel sendThread recvThread

          pure ()


  case config ^. role of
    Local -> do
      _ <- localThread
      pure ()
    Remote -> do
      _ <- remoteThread
      pure ()

  pure ()