{-# LANGUAGE ScopedTypeVariables #-}

import System.Environment
import System.Process
import Control.Monad
import Control.Concurrent
import Control.Exception
import qualified Control.Exception as CE
import Network.Socket
import System.Random
import System.IO
import qualified Data.ByteString as BS
import Text.Printf (printf)

main :: IO ()
main = do
  listenPort' : targetPort' : sshArgs <- getArgs
  let listenPort = read listenPort' :: Int
      targetPort = read targetPort' :: Int
  sock <- socket AF_INET Stream defaultProtocol
  setSocketOption sock ReuseAddr 1
  addr <- inet_addr "127.0.0.1"
  bindSocket sock $ SockAddrInet (fromIntegral listenPort) addr
  listen sock 5
  putStrLn $ "listening for connections on port " ++ show listenPort
  forever $ do
    (conn, connaddr) <- accept sock
    hConn <- socketToHandle conn ReadWriteMode
    putStrLn $ "connection accepted: " ++ show connaddr
    forkIO $ do
      tunnelPort :: Int <- randomRIO (50000, 55000)
      handleConnection hConn tunnelPort targetPort sshArgs

handleConnection :: Handle -> Int -> Int -> [String] -> IO ()
handleConnection hConn tunnelPort targetPort sshArgs = do
    (_, _, _, p) <- createProcess (proc "ssh" $ ["-v"] ++ [
      "-L" ++ show tunnelPort ++ ":127.0.0.1:" ++ show targetPort, "-n"] ++
      sshArgs)
    finally forwardToTunnel $ do
      putStrLn "terminating ssh"
      terminateProcess p
      _ <- waitForProcess p
      putStrLn "ssh exited"
  
  where forwardToTunnel = do
          hTunn <- tryTunnelConnection tunnelPort 15
          _ <- forkIO $ forward "hTunn->hConn" hTunn hConn
          forward "hConn->hTunn" hConn hTunn
          
        forward desc h1 h2 = do
          forever $ do
            bs <- BS.hGetSome h1 4096
            putStrLn $ printf "%s %d bytes" desc (BS.length bs)
            if BS.null bs then do
                error "no more data to tunnel"
            else do
                BS.hPut h2 bs
                hFlush h2

tryTunnelConnection :: Int -> Int -> IO Handle
tryTunnelConnection _    0        = error "can't connect to tunnel"
tryTunnelConnection port attempts = do
  threadDelay (1*1000*1000)
  tunn <- socket AF_INET Stream defaultProtocol
  addr <- inet_addr "127.0.0.1"
  CE.catch (do connect tunn $ SockAddrInet (fromIntegral port) addr
               hTunn <- socketToHandle tunn ReadWriteMode
               putStrLn "got hTunn" 
               return hTunn)
        (\(_::SomeException) -> do
            sClose tunn
            tryTunnelConnection port (attempts-1))