summaryrefslogtreecommitdiff
path: root/dht/examples/consolation.hs
blob: 0c576dfc7a46f684013f27553c3005bdc51532fa (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
{-# LANGUAGE OverloadedStrings #-}
module Main where

import Control.Monad
import Control.Applicative
import Control.Concurrent
import Control.Concurrent.STM
import Data.Monoid
import Data.Char
import System.INotify ( initINotify, EventVariety(Modify), addWatch )
import Data.Word      ( Word8 )
import Data.Text      ( Text )
import Data.Map       ( Map )
import Data.List      ( foldl' )
import qualified Data.Map as Map
import qualified Data.Traversable as Traversable
import qualified Data.Text as Text
import qualified Data.Text.IO as Text
import qualified Network.BSD as BSD

import WaitForSignal  ( waitForTermSignal )
import UTmp           ( users2, utmp_file, UtmpRecord(..), UT_Type(..) )
import FGConsole      ( monitorTTY )

data ConsoleState = ConsoleState
    { csActiveTTY :: TVar Word8
    , csUtmp :: TVar (Map Text (TVar (Maybe UtmpRecord)))
    }

newConsoleState = atomically $
    ConsoleState <$> newTVar 0 <*> newTVar Map.empty


onLogin cs start = \e -> do
    us <- UTmp.users2
    let (m,cruft) =
            foldl' (\(m,cruft) x ->
                        case utmpType x of
                         USER_PROCESS
                           -> (Map.insert (utmpTty x) x m,cruft)
                         DEAD_PROCESS | utmpPid x /= 0
                           -> (m,Map.insert (utmpTty x) x cruft)
                         _ -> (m,cruft))
                   (Map.empty,Map.empty)
                   us
    forM_ (Map.elems cruft) $ \c -> do
        putStrLn $ "cruft " ++ show (utmpTty c, utmpPid c,utmpHost c, utmpRemoteAddr c)
    newborn <- atomically $ do
        old <- readTVar (csUtmp cs) -- swapTVar (csUtmp cs) m
        newborn <- flip Traversable.mapM (m Map.\\ old)
                    $ newTVar . Just
        updated <- let upd v u = writeTVar v $ Just u
                   in Traversable.sequence $ Map.intersectionWith upd old m
        let dead = old Map.\\ m
        Traversable.mapM (flip writeTVar Nothing) dead
        writeTVar (csUtmp cs) $ (old `Map.union` newborn) Map.\\ dead
        return newborn
    let getActive = do
            tty <- readTVar $ csActiveTTY cs
            utmp <- readTVar $ csUtmp cs
            flip (maybe $ return (tty,Nothing))
                 (Map.lookup ("tty"<>tshow tty) utmp)
                $ \tuvar -> do
            tu <- readTVar tuvar
            return (tty,tu)
            
    forM_ (Map.elems newborn) $ 
        forkIO . start getActive
    -- forM_ (Map.elems dead   ) $ putStrLn . ("gone: "++) . show

onTTY outvar cs vtnum = do
    logit outvar $ "switch: " <> tshow vtnum
    atomically $ writeTVar (csActiveTTY cs) vtnum

retryWhen var pred = do
    value <- var
    if pred value then retry
                  else return value

tshow x = Text.pack . show $ x

resource :: UtmpRecord -> Text
resource u =
    case utmpTty u of
        s | Text.take 3 s == "tty" -> s
        s | Text.take 4 s == "pts/" -> "pty" <> Text.drop 4 s <> ":" <> utmpHost u
        s -> escapeR s <> ":" <> utmpHost u
 where
    escapeR s = s

textHostName = fmap Text.pack BSD.getHostName

ujid u = do
    h <- textHostName
    return $ utmpUser u <> "@" <> h <> "/" <> resource u

newCon :: (Text -> IO ()) -> STM (Word8,Maybe UtmpRecord) -> TVar (Maybe UtmpRecord) -> IO ()
newCon log activeTTY utmp = do
    ((tty,tu),u) <- atomically $
        liftM2 (,) activeTTY
                   (readTVar utmp)
    flip (maybe $ return ()) u $ \u -> do
    jid <- ujid u
    log $ status (resource u) tty tu <> " " <> jid <> "  pid=" <> tshow (utmpPid u)
            <> (if istty (resource u)
                 then " host=" <> tshow (utmpHost u)
                 else "")
            <> " session=" <> tshow (utmpSession u)
            <> " addr=" <> tshow (utmpRemoteAddr u)
    loop tty tu (Just u)
 where
    bstatus r ttynum mtu
        =  r == ttystr
        || match mtu
        where ttystr = "tty" <> tshow ttynum
              searchstr mtu = maybe ttystr utmpHost $ do
                tu <- mtu
                guard (not $ Text.null $ utmpHost tu)
                return tu
              match mtu = searchstr mtu `Text.isInfixOf` Text.dropWhile (/=':') r
    status r ttynum tu =
        if bstatus r ttynum tu
            then "Available"
            else "Away     "

    istty r = fst3 == "tty" && Text.all isDigit rst
     where
        (fst3,rst) = Text.splitAt 3 r

    loop tty tu u = do
        what <- atomically $ foldr1 orElse
            [ do (tty',tu') <- retryWhen activeTTY
                            (\ttyu -> bstatus r tty tu == uncurry (bstatus r) ttyu)
                 return $ ttyChanged tty' tu'
            , do u' <- retryWhen (readTVar utmp) (==u)
                 return $ utmpChanged u'
            ]
        what
      where
        r = maybe "" resource u

        ttyChanged tty' tu' = do
            jid <- maybe (return "") ujid u
            log $ status r tty' tu' <> " " <> jid
            loop tty' tu' u

        utmpChanged u' = maybe dead changed u'
         where
            changed u' = do
                jid0 <- maybe (return "") ujid u
                jid <- ujid u'
                log $ "changed: " <> jid0 <> " --> " <> jid
                loop tty tu (Just u')
            dead = do
                jid <- maybe (return "") ujid u
                log $ "Offline   " <> jid

logit outvar s = do
    atomically $ takeTMVar outvar
    Text.putStrLn s
    atomically $ putTMVar outvar ()


main = do
    outvar <- atomically $ newTMVar ()

    cs <- newConsoleState
    inotify <- initINotify

    -- get active tty
    mtty <- monitorTTY (onTTY outvar cs)
    atomically $ retryWhen (readTVar $ csActiveTTY cs) (==0)

    -- read utmp
    onLogin cs (newCon $ logit outvar) Modify

    -- monitor utmp
    wd <- addWatch
            inotify
            [Modify] -- [CloseWrite,Open,Close,Access,Modify,Move]
            utmp_file
            (onLogin cs (newCon $ logit outvar))

    waitForTermSignal

    putStrLn "goodbye."