tls-debug-0.2.5: src/Stunnel.hs

{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE ScopedTypeVariables #-}
import Network.BSD
import Network.Socket
import System.IO
import System.IO.Error (isEOFError)
import System.Console.CmdArgs
import System.Certificate.X509

import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as L

import Control.Concurrent (forkIO)
import Control.Concurrent.MVar
import Control.Exception (finally, throw, SomeException)
import qualified Control.Exception as E
import Control.Monad (when, forever)

import Data.Char (isDigit)

import qualified Crypto.Random.AESCtr as RNG
import Network.TLS
import Network.TLS.Extra

ciphers :: [Cipher]
ciphers =
    [ cipher_AES128_SHA1
    , cipher_AES256_SHA1
    , cipher_RC4_128_MD5
    , cipher_RC4_128_SHA1
    ]

loopUntil :: Monad m => m Bool -> m ()
loopUntil f = f >>= \v -> if v then return () else loopUntil f

readOne h = do
    r <- E.try $ hWaitForInput h (-1)
    case r of
        Left err    -> if isEOFError err then return B.empty else throw err
        Right True  -> B.hGetNonBlocking h 4096
        Right False -> return B.empty

tlsclient :: Handle -> TLSCtx -> IO ()
tlsclient srchandle dsthandle = do
    hSetBuffering srchandle NoBuffering

    handshake dsthandle

    _ <- forkIO $ forever $ do
        dat <- recvData dsthandle
        putStrLn ("received " ++ show dat)
        B.hPut srchandle dat
    loopUntil $ do
        b <- readOne srchandle
        putStrLn ("sending " ++ show b)
        if B.null b
            then do
                bye dsthandle
                return True
            else do
                sendData dsthandle (L.fromChunks [b])
                return False
    return ()

tlsserver srchandle dsthandle = do
    hSetBuffering dsthandle NoBuffering

    handshake srchandle

    loopUntil $ do
        d <- recvData srchandle
        putStrLn ("received: " ++ show d)
        sendData srchandle (L.pack $ map (toEnum . fromEnum) "this is some data")
        return False
    putStrLn "end"

data MemSessionManager = MemSessionManager (MVar [(SessionID, SessionData)])

instance SessionManager MemSessionManager where
    sessionEstablish (MemSessionManager mvar) sid sdata = modifyMVar_ mvar (\l -> return $ (sid,sdata) : l)
    sessionResume (MemSessionManager mvar) sid          = withMVar mvar (return . lookup sid)
    sessionInvalidate (MemSessionManager mvar) _        = return ()

clientProcess certs handle dsthandle dbg sessionStorage _ = do
    rng <- RNG.makeSystem
    let logging = if not dbg
            then defaultLogging
            else defaultLogging { loggingPacketSent = putStrLn . ("debug: send: " ++)
                                , loggingPacketRecv = putStrLn . ("debug: recv: " ++)
                                }

    let serverstate = maybe id (setSessionManager . MemSessionManager) sessionStorage $ defaultParamsServer
                        { pAllowedVersions = [SSL3,TLS10,TLS11,TLS12]
                        , pCiphers         = ciphers
                        , pCertificates    = certs
                        , pLogging         = logging
                        }

    ctx <- contextNewOnHandle handle serverstate rng
    tlsserver ctx dsthandle

data Stunnel =
      ClientConfig
        { destinationType :: String
        , destination     :: String
        , sourceType      :: String
        , source          :: String
        , debug           :: Bool
        , validCert       :: Bool }
    | ServerConfig
        { destinationType :: String
        , destination     :: String
        , sourceType      :: String
        , source          :: String
        , debug           :: Bool
        , disableSession  :: Bool
        , certificate     :: FilePath
        , key             :: FilePath }
    deriving (Show, Data, Typeable)

clientOpts = ClientConfig
    { destinationType = "tcp"             &= help "type of source (tcp, unix, fd)" &= typ "DESTTYPE"
    , destination     = "localhost:6061"  &= help "destination address influenced by destination type" &= typ "ADDRESS"
    , sourceType      = "tcp"             &= help "type of source (tcp, unix, fd)" &= typ "SOURCETYPE"
    , source          = "localhost:6060"  &= help "source address influenced by source type" &= typ "ADDRESS"
    , debug           = False             &= help "debug the TLS protocol printing debugging to stdout" &= typ "Bool"
    , validCert       = False             &= help "check if the certificate receive is valid" &= typ "Bool"
    }
    &= help "connect to a remote destination that use SSL/TLS"
    &= name "client"

serverOpts = ServerConfig
    { destinationType = "tcp"             &= help "type of source (tcp, unix, fd)" &= typ "DESTTYPE"
    , destination     = "localhost:6060"  &= help "destination address influenced by destination type" &= typ "ADDRESS"
    , sourceType      = "tcp"             &= help "type of source (tcp, unix, fd)" &= typ "SOURCETYPE"
    , source          = "localhost:6061"  &= help "source address influenced by source type" &= typ "ADDRESS"
    , disableSession  = False             &= help "disable support for session" &= typ "Bool"
    , debug           = False             &= help "debug the TLS protocol printing debugging to stdout" &= typ "Bool"
    , certificate     = "certificate.pem" &= help "X509 public certificate to use" &= typ "FILE"
    , key             = "certificate.key" &= help "private key linked to the certificate" &= typ "FILE"
    }
    &= help "listen for connection that use SSL/TLS and relay it to a different connection"
    &= name "server"

mode = cmdArgsMode $ modes [clientOpts,serverOpts]
    &= help "create SSL/TLS tunnel in client or server mode" &= program "stunnel" &= summary "Stunnel v0.1 (Haskell TLS)"

data StunnelAddr   =
      AddrSocket Family SockAddr
    | AddrFD Handle Handle

data StunnelHandle =
      StunnelSocket Socket
    | StunnelFd     Handle Handle

getAddressDescription :: String -> String -> IO StunnelAddr
getAddressDescription "tcp"  desc = do
    let (s, p) = break ((==) ':') desc
    when (p == "") (error "missing port: expecting [source]:port")
    pn <- if and $ map isDigit $ drop 1 p
        then return $ fromIntegral $ (read (drop 1 p) :: Int)
        else do
            service <- getServiceByName (drop 1 p) "tcp"
            return $ servicePort service
    he <- getHostByName s
    return $ AddrSocket AF_INET (SockAddrInet pn (head $ hostAddresses he))

getAddressDescription "unix" desc = do
    return $ AddrSocket AF_UNIX (SockAddrUnix desc)

getAddressDescription "fd" _  =
    return $ AddrFD stdin stdout

getAddressDescription _ _  = error "unrecognized source type (expecting tcp/unix/fd)"

connectAddressDescription (AddrSocket family sockaddr) = do
    sock <- socket family Stream defaultProtocol
    E.catch (connect sock sockaddr)
          (\(e :: SomeException) -> sClose sock >> error ("cannot open socket " ++ show sockaddr ++ " " ++ show e))
    return $ StunnelSocket sock

connectAddressDescription (AddrFD h1 h2) = do
    return $ StunnelFd h1 h2

listenAddressDescription (AddrSocket family sockaddr) = do
    sock <- socket family Stream defaultProtocol
    E.catch (bindSocket sock sockaddr >> listen sock 10 >> setSocketOption sock ReuseAddr 1)
          (\(e :: SomeException) -> sClose sock >> error ("cannot open socket " ++ show sockaddr ++ " " ++ show e))
    return $ StunnelSocket sock

listenAddressDescription (AddrFD _ _) = do
    error "cannot listen on fd"

doClient :: Stunnel -> IO ()
doClient pargs = do
    srcaddr <- getAddressDescription (sourceType pargs) (source pargs)
    dstaddr <- getAddressDescription (destinationType pargs) (destination pargs)

    let logging = if not $ debug pargs then defaultLogging else defaultLogging
                            { loggingPacketSent = putStrLn . ("debug: send: " ++)
                            , loggingPacketRecv = putStrLn . ("debug: recv: " ++)
                            }

    store <- getSystemCertificateStore
    let crecv = if validCert pargs then certificateVerifyChain store else (\_ -> return CertificateUsageAccept)
    let clientstate = defaultParamsClient { pConnectVersion = TLS10
                                          , pAllowedVersions = [TLS10,TLS11,TLS12]
                                          , pCiphers = ciphers
                                          , pCertificates = []
                                          , pLogging = logging
                                          , onCertificatesRecv = crecv
                                          }

    case srcaddr of
        AddrSocket _ _ -> do
            (StunnelSocket srcsocket) <- listenAddressDescription srcaddr
            forever $ do
                (s, _) <- accept srcsocket
                rng    <- RNG.makeSystem
                srch   <- socketToHandle s ReadWriteMode

                (StunnelSocket dst)  <- connectAddressDescription dstaddr

                dsth <- socketToHandle dst ReadWriteMode
                dstctx <- contextNewOnHandle dsth clientstate rng
                _    <- forkIO $ finally
                    (tlsclient srch dstctx)
                    (hClose srch >> hClose dsth)
                return ()
        AddrFD _ _ -> error "bad error fd. not implemented"

doServer :: Stunnel -> IO ()
doServer pargs = do
    cert    <- fileReadCertificate $ certificate pargs
    pk      <- fileReadPrivateKey $ key pargs
    srcaddr <- getAddressDescription (sourceType pargs) (source pargs)
    dstaddr <- getAddressDescription (destinationType pargs) (destination pargs)

    sessionStorage <- if disableSession pargs then return Nothing else (Just `fmap` newMVar [])

    case srcaddr of
        AddrSocket _ _ -> do
            (StunnelSocket srcsocket) <- listenAddressDescription srcaddr
            forever $ do
                (s, addr) <- accept srcsocket
                srch <- socketToHandle s ReadWriteMode
                r <- connectAddressDescription dstaddr
                dsth <- case r of
                    StunnelFd _ _     -> return stdout
                    StunnelSocket dst -> socketToHandle dst ReadWriteMode

                _ <- forkIO $ finally
                    (clientProcess [(cert, Just pk)] srch dsth (debug pargs) sessionStorage addr >> return ())
                    (hClose srch >> (when (dsth /= stdout) $ hClose dsth))
                return ()
        AddrFD _ _ -> error "bad error fd. not implemented"

main :: IO ()
main = do
    x <- cmdArgsRun mode
    case x of
        ClientConfig {} -> doClient x
        ServerConfig {} -> doServer x