OnRmt-1.0.0.0: onrmtssh.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeFamilies #-}

module Main where

import           Brick.Widgets.ProgressBar (progressCompleteAttr)
import           Data.List (partition, mapAccumL)
import           Data.Monoid
import qualified Data.Text as T
import Graphics.Vty.Attributes ( defAttr, withStyle
                               , withForeColor, withBackColor
                               , magenta, cyan
                               )
import           Network.OnRmt
import           Network.OnRmt.UI.BrickUI
import           Network.SSH.KnownHosts (SSHRemote, readKnownHosts
                                        , targetAddr
                                        , isIPv4Remote, isIPv6Remote
                                        , isDNSRemote
                                        )
import           System.Environment


rmtTestAttrs =
    [ (progressCompleteAttr
      , defAttr `withForeColor` magenta `withBackColor` cyan)
    ]


params = OnRmtParams { appName = "onrmttest"
                     , maxParallel = 9
                     , directSSH = False }


data SSHGroups = SSHGroups { ipv4 :: [SSHRemote]
                           , ipv6 :: [SSHRemote]
                           , dns :: [SSHRemote]
                           , other :: [SSHRemote]
                           }


getEntries =
  do allknown <- readKnownHosts
     let (ipv4',not1) = partition isIPv4Remote allknown
         (ipv6',not2) = partition isIPv6Remote not1
         (dns',other') = partition isDNSRemote not2
     return $ SSHGroups ipv4' ipv6' dns' other'


remoteWorkItems :: SSHGroups -> [WorkItems]
remoteWorkItems sshg =
    [ WorkGroup "IPv4" (WorkItems $ length $ ipv4 sshg)
    , WorkGroup "IPv6" (WorkItems $ length $ ipv6 sshg)
    , WorkGroup "DNS" (WorkItems $ length $ dns sshg)
    , WorkGroup "Other" (WorkItems $ length $ other sshg)
    ]

instance WorkEntry SSHRemote where
    name = fst . targetAddr
    identify r _ = name r
    rmtaddr = name

instance WorkGroup SSHGroups where
    type GroupEntry SSHGroups = SSHRemote
    getItems = remoteWorkItems
    numEntries g = sum [ length $ ipv4 g
                       , length $ ipv6 g
                       , length $ dns g
                       , length $ other g
                       ]
    getEntry g n =
        let l1 = length $ ipv4 g
            l2 = length $ ipv6 g
            l3 = length $ dns g
            l4 = length $ other g
        in if n < l1
           then ipv4 g !! n
           else if n < (l1 + l2)
                then ipv6 g !! (n - l1)
                else if n < (l1 + l2 + l3)
                     then dns g !! (n - l1 - l2)
                     else other g !! (n - l1 - l2 - l3)

main =
    do (uiMode, stateGen, s) <- brickUI (appName params) rmtTestAttrs
       result <- onRmt getEntries params uiMode stateGen s
       putStr "Final status: "
       putStrLn $ show result
       return ()