{-# 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 qualified Data.ByteString.Lazy.Char8 as L 
    ( unpack
    , pack
    , take
    , putStrLn
    )
import Data.List (groupBy)
import Data.Maybe (listToMaybe,mapMaybe,catMaybes)

import Network.Socket
import System.Environment
import Control.Arrow (first)
import System.Endian

usage = do
    putStrLn $ "whosocket numeric-address port"

main = do
    args <- getArgs
    case (args??0,args??1) of
        (Just addr_str,Just port_str) -> whosocket addr_str port_str
        _                             -> usage

whosocket addr_str port_str = do
    info <- getAddrInfo (Just $ defaultHints { addrFlags = [ AI_NUMERICHOST ] }) 
                        (Just addr_str) 
                        (Just port_str)
    let addr = head $ map addrAddress info
    r <- getLocalPeerCred' addr
    putStrLn $ "r{"++show addr++"} = " ++ show r

    us <- UTmp.users
    let filterTTYs (_,tty,pid) = 
            if L.take 3 tty == "tty"
              then Just (tty,pid)
              else Nothing
        tty_pids = mapMaybe filterTTYs us

    tty <- maybe (return Nothing)
                 (uncurry $ identifyTTY tty_pids)
                 r
    putStrLn $ "uid = " ++ show (fmap fst r)
    L.putStrLn $ "tty = " <++?> tty

    return ()


makeUidStr "4294967295" = "invalid"
makeUidStr uid          = uid


scanProc uid inode = do
    contents <- getDirectoryContents "/proc" `catchIO_` return []
    let pids = reverse $ filter (\n -> not (null n) && isDigit (head n)) contents
    let searchPids []         = return Nothing
        searchPids (pid:pids) =  do
            loginuid <- fmap makeUidStr $ readFile $ "/proc/"++pid++"/loginuid"
            if False -- (uid/=loginuid) -- this check proved bad when mcabber ran on tty3
             then searchPids pids
             else do
                -- putStrLn $ "pid "++show pid ++ " --> uid "++show loginuid
                let loop [] = return Nothing
                    loop ("0":fds) = loop fds
                    loop (fd:fds) = do
                      handleIO_ (loop fds) $ do
                          what <- readSymbolicLink $ "/proc/"++pid++"/fd/"++fd
                          -- putStrLn $ " what= "++show what
                          if (what=="socket:["++inode++"]")
                            then do
                              tty <- readSymbolicLink $ "/proc/"++pid++"/fd/0"
                              return (Just (pid,tty))
                            else loop fds
                fds <- getDirectoryContents ("/proc/"++pid++"/fd") `catchIO_` return []
                mb <- loop fds
                maybe (searchPids pids) (return . Just) mb

    fmap (fmap (first (read :: String -> CPid))) $ searchPids pids

ttyOrDisplay pid ttydev = do
    ptty <- searchParentsForTTY (show pid) ttydev
    case ptty of
        Just tty -> return tty
        Nothing  -> do
            display <- readDisplayVariable pid
            -- putStrLn $ "display = " ++ show display
            case display of
                Just (_,disp) -> return disp
                _             -> return ttydev

searchParentsForTTY pid  ttydev | take 8 ttydev == "/dev/tty" = return . Just $ drop 5 ttydev
searchParentsForTTY "1"  ttydev | otherwise                   = return Nothing
searchParentsForTTY pid ttydev = do
    stat <- handleIO_ (return "") . readFile $ "/proc/"++pid++"/stat"
    case words stat ?? 3 of
        Nothing   -> return Nothing
        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 tty_pids 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
    putStrLn $ "users = " ++ show tty_pids
    dts <- ttyToXorgs tty_pids

    putStrLn $ "displays = " ++ show dts 
    putStrLn $ "tty = " ++ show tty
    -- -- displays = [(":5",Chunk "tty7" Empty)]

    let tty' = if take 3 tty=="tty" then Just (L.pack tty) else lookup (parseTty tty) (map (first parseTty) dts)
    return tty'
    where
        parseTty :: String -> Float
        parseTty = read .  tail . dropWhile (/=':')