1 files changed, 70 insertions, 20 deletions

diff --git a/whosocket.hs b/whosocket.hs
index 6955830a..420b707a 100644
--- a/whosocket.hs
+++ b/whosocket.hs
@@ -1,15 +1,24 @@
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE TupleSections     #-}
 module Main where
 
 import LocalPeerCred
 import ControlMaybe
+import UTmp
+import ByteStringOperators
 
 import System.Directory
 import Data.Char
 import System.Posix.Types
 import System.Posix.Files
-import Data.ByteString.Lazy.Char8 as L (unpack)
+import qualified Data.ByteString.Lazy.Char8 as L 
+    ( unpack
+    , pack
+    , take
+    , putStrLn
+    )
 import Data.List (groupBy)
-import Data.Maybe (listToMaybe)
+import Data.Maybe (listToMaybe,mapMaybe,catMaybes)
 
 import Network.Socket
 import System.Environment
@@ -30,15 +39,13 @@ whosocket addr_str port_str = do
                         (Just port_str)
     let addr = head $ map addrAddress info
     r <- getLocalPeerCred' addr
-    putStrLn $ "r{"++show addr++"} = " ++ show r
+    -- putStrLn $ "r{"++show addr++"} = " ++ show r
 
-    let Just (uid,inode) = r
-    pid <- scanProc (show uid) (L.unpack inode)
-    putStrLn $ "scanProc --> "++show pid
-    withJust pid $ \(pid,ttydev) -> do
-    tty <- ttyOrDisplay pid ttydev
-    putStrLn $ "pid = " ++ show pid
-    putStrLn $ "tty = " ++ show tty
+    tty <- maybe (return Nothing)
+                 (uncurry identifyTTY)
+                 r
+    putStrLn $ "uid = " ++ show (fmap fst r)
+    L.putStrLn $ "tty = " <++?> tty
 
     return ()
 
@@ -79,16 +86,8 @@ ttyOrDisplay pid ttydev = do
     case ptty of
         Just tty -> return tty
         Nothing  -> do
-            env <- handleIO_ (return "") . readFile $ "/proc/"++show pid++"/environ"
-            let vs = unzero $ groupBy (\_ c->c/='\0') env
-                unzero []     = []
-                unzero (v:vs) = v:map tail vs
-                keyvalue xs = (key,value) 
-                  where
-                    (key,ys) = break (=='=') xs
-                    value = case ys of { [] -> [];  (_:ys') -> ys' }
-                display = listToMaybe . filter ((=="DISPLAY").fst) . map keyvalue $ vs
-            putStrLn $ "display = " ++ show display
+            display <- readDisplayVariable pid
+            -- putStrLn $ "display = " ++ show display
             case display of
                 Just (_,disp) -> return disp
                 _             -> return ttydev
@@ -102,3 +101,54 @@ searchParentsForTTY pid ttydev = do
         Just ppid -> do
           tty <- handleIO_ (return "") $ readSymbolicLink $ "/proc/"++ppid++"/fd/0"
           searchParentsForTTY ppid tty
+
+readDisplayVariable pid = do
+            env <- handleIO_ (return "") 
+                    . readFile $ "/proc/"++show pid++"/environ"
+            let vs = unzero $ groupBy (\_ c->c/='\0') env
+                unzero []     = []
+                unzero (v:vs) = v:map tail vs
+                keyvalue xs = (key,value) 
+                  where
+                    (key,ys) = break (=='=') xs
+                    value = case ys of { [] -> [];  (_:ys') -> ys' }
+                display = listToMaybe 
+                            . filter ((=="DISPLAY").fst) 
+                            . map keyvalue 
+                            $ vs
+            return display 
+
+ttyToXorgs tty_pids = do
+    dts' <- flip mapM tty_pids $ \(tty,pid) -> do
+        cmd' <- readFile $ "/proc/"++show pid++"/cmdline"
+        case listToMaybe . words . takeWhile (/='\0') $ cmd' of
+          Nothing  -> return Nothing
+          Just cmd -> do
+            if notElem cmd ["gdm-session-worker"]
+                then return Nothing
+                else do
+                  display <- readDisplayVariable pid
+                  return (fmap ( (,tty) . snd ) display)
+    let dts = catMaybes dts'
+    return dts
+
+identifyTTY uid inode = do
+    pid <- scanProc (show uid) (L.unpack inode)
+    -- putStrLn $ "scanProc --> "++show pid
+    flip (maybe (return Nothing)) pid $ \(pid,ttydev) -> do
+    tty <- ttyOrDisplay pid ttydev
+
+    us <- users
+    let filterTTYs (_,tty,pid) = 
+            if L.take 3 tty == "tty"
+              then Just (tty,pid)
+              else Nothing
+        tty_pids = mapMaybe filterTTYs us
+    -- putStrLn $ "users = " ++ show tty_pids
+    dts <- ttyToXorgs tty_pids
+
+    -- putStrLn $ "displays = " ++ show dts 
+    -- -- displays = [(":5",Chunk "tty7" Empty)]
+
+    let tty' = if take 3 tty=="tty" then Just (L.pack tty) else lookup tty dts
+    return tty'