A basic application written using ssh-hans

simple_hans_server.hs

module Main where

{-
Here is a real---if not particularly interesting---interaction with this server:

  [gdritter@armilla ~]$ ssh -p 9999 localhost incr
  ok.
  [gdritter@armilla ~]$ ssh -p 9999 localhost incr
  ok.
  [gdritter@armilla ~]$ ssh -p 9999 localhost double
  ok.
  [gdritter@armilla ~]$ ssh -p 9999 localhost get
  4
  [gdritter@armilla ~]$ ssh -p 9999 localhost decr
  ok.
  [gdritter@armilla ~]$ ssh -p 9999 localhost add 22
  ok.
  [gdritter@armilla ~]$ ssh -p 9999 localhost get
  25
-}

import           Control.Concurrent (forkIO)
import           Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import           Data.ByteString.Char8 (pack, unpack)
import           Data.ByteString.Lazy (toStrict)
import           Data.IORef (IORef, newIORef, modifyIORef, readIORef)
import           Data.Monoid ((<>))
import           Network ( PortID(..)
                         , Socket
                         , listenOn
                         , accept
                         , withSocketsDo
                         )
import           Network.SSH.Server ( Server(..)
                                    , ServerCredential
                                    , SessionHandlers(..)
                                    , HandleLike
                                    , AuthResult(AuthAccepted)
                                    , loadPrivateKeys
                                    , handle2HandleLike
                                    , sshServer
                                    )

-- we hard-code in the port number and location of the private keys
main :: IO ()
main = withSocketsDo $ do
  sock <- listenOn (PortNumber 9999)
  keys <- loadPrivateKeys "server-keys"
  ref  <- newIORef 0
  sshServer (server ref sock keys)

-- Our SSH server structure. the sDebugLevel should be non-zero if we want richer
-- debug output
server :: IORef Int -> Socket -> [ServerCredential] -> Server
server ref sock keys = Server
  { sAccept             = handleConnection ref sock
  , sAuthenticationAlgs = keys
  , sVersion            = "SSH_Hans_Test_0.1"
  , sDebugLevel         = 0
  }

-- This accepts every connection for ease of use, and doesn't accept any kind
-- of connection other than executing commands: it will refuse any request
-- to create a shell, for example. The commands in question are described below,
-- but all are simple textual commands that manipulate a stored integer value.
handleConnection :: IORef Int -> Socket -> IO (SessionHandlers, HandleLike)
handleConnection ref sock = do
  (handle, _, _) <- accept sock
  let handlers = SessionHandlers
        { cOpenShell        = \ _ _ _ _ _ -> return False
        , cDirectTcp        = \ _ _ _ _   -> return False
        , cRequestSubsystem = \ _ _ _     -> return False
        , cRequestExec      = \ cmd _ write -> do
             _ <- forkIO (runCmd ref cmd write)
             return True
        , cAuthHandler      = \ _ _ _ _ -> do
             return AuthAccepted
        }
  return (handlers, handle2HandleLike handle)

-- Our command language manipulated a single integer value in small
-- ways, and can result in a value being printed.
runCmd :: IORef Int -> ByteString -> (Maybe ByteString -> IO ()) -> IO ()
runCmd ref cmd write = case BS.split 32 cmd of
    ["incr"]     -> runAction (+1)
    ["double"]   -> runAction (*2)
    ["decr"]     -> runAction (\ x -> x - 1)
    ["add", val] ->
      let num = read (unpack val)
      in runAction (+num)
    ["get"] -> do
      n <- readIORef ref
      respond (pack (show n))
    _ -> respond usage
  where respond msg = write (Just (msg <> "\n")) >> write Nothing
        runAction f = modifyIORef ref f >> respond "ok."

-- The usage string is printed if the command given is not valid.
usage :: ByteString
usage = "Command must be one of:\n incr\n decr\n double\n add [num]\n get\n"