{-# LANGUAGE CPP #-}
{-# LANGUAGE OverloadedStrings #-}

module Program.Mighty.Process (
    getMightyPid
  ) where

import Control.Monad.Trans.Resource (runResourceT)
import Data.ByteString.Char8 (ByteString)
import qualified Data.ByteString.Char8 as BS
import Data.Conduit
import qualified Data.Conduit.Binary as CB
import qualified Data.Conduit.List as CL
import Data.Conduit.Process
import Data.Function
import Data.List
import System.Posix.Types

----------------------------------------------------------------

data PsResult = PsResult {
    PsResult -> ByteString
uid :: ByteString
  , PsResult -> ProcessID
pid :: ProcessID
  , PsResult -> ProcessID
ppid :: ProcessID
  , PsResult -> ByteString
command :: ByteString
  } deriving (PsResult -> PsResult -> Bool
(PsResult -> PsResult -> Bool)
-> (PsResult -> PsResult -> Bool) -> Eq PsResult
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: PsResult -> PsResult -> Bool
$c/= :: PsResult -> PsResult -> Bool
== :: PsResult -> PsResult -> Bool
$c== :: PsResult -> PsResult -> Bool
Eq, Int -> PsResult -> ShowS
[PsResult] -> ShowS
PsResult -> String
(Int -> PsResult -> ShowS)
-> (PsResult -> String) -> ([PsResult] -> ShowS) -> Show PsResult
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PsResult] -> ShowS
$cshowList :: [PsResult] -> ShowS
show :: PsResult -> String
$cshow :: PsResult -> String
showsPrec :: Int -> PsResult -> ShowS
$cshowsPrec :: Int -> PsResult -> ShowS
Show)

toPsResult :: [ByteString] -> PsResult
toPsResult :: [ByteString] -> PsResult
toPsResult (ByteString
a:ByteString
b:ByteString
c:ByteString
_:ByteString
_:ByteString
_:ByteString
_:ByteString
h:[ByteString]
_) = PsResult :: ByteString -> ProcessID -> ProcessID -> ByteString -> PsResult
PsResult {
    uid :: ByteString
uid     = ByteString
a
  , pid :: ProcessID
pid     = ProcessID
-> ((Int, ByteString) -> ProcessID)
-> Maybe (Int, ByteString)
-> ProcessID
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ProcessID
0 (Int -> ProcessID
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> ProcessID)
-> ((Int, ByteString) -> Int) -> (Int, ByteString) -> ProcessID
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, ByteString) -> Int
forall a b. (a, b) -> a
fst) (Maybe (Int, ByteString) -> ProcessID)
-> Maybe (Int, ByteString) -> ProcessID
forall a b. (a -> b) -> a -> b
$ ByteString -> Maybe (Int, ByteString)
BS.readInt ByteString
b
  , ppid :: ProcessID
ppid    = ProcessID
-> ((Int, ByteString) -> ProcessID)
-> Maybe (Int, ByteString)
-> ProcessID
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ProcessID
0 (Int -> ProcessID
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> ProcessID)
-> ((Int, ByteString) -> Int) -> (Int, ByteString) -> ProcessID
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, ByteString) -> Int
forall a b. (a, b) -> a
fst) (Maybe (Int, ByteString) -> ProcessID)
-> Maybe (Int, ByteString) -> ProcessID
forall a b. (a -> b) -> a -> b
$ ByteString -> Maybe (Int, ByteString)
BS.readInt ByteString
c
  , command :: ByteString
command = ByteString
h
  }
toPsResult [ByteString]
_ = ByteString -> ProcessID -> ProcessID -> ByteString -> PsResult
PsResult ByteString
"unknown" ProcessID
0 ProcessID
0 ByteString
"unknown"

----------------------------------------------------------------

runPS :: IO [PsResult]
runPS :: IO [PsResult]
runPS = (ExitCode, [PsResult]) -> [PsResult]
forall a b. (a, b) -> b
snd ((ExitCode, [PsResult]) -> [PsResult])
-> IO (ExitCode, [PsResult]) -> IO [PsResult]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ResourceT IO (ExitCode, [PsResult]) -> IO (ExitCode, [PsResult])
forall (m :: * -> *) a. MonadUnliftIO m => ResourceT m a -> m a
runResourceT (String
-> ConduitT ByteString Void (ResourceT IO) [PsResult]
-> ResourceT IO (ExitCode, [PsResult])
forall (m :: * -> *) a.
MonadIO m =>
String -> ConduitT ByteString Void m a -> m (ExitCode, a)
sourceCmdWithConsumer String
"ps -ef" ConduitT ByteString Void (ResourceT IO) [PsResult]
forall c. ConduitM ByteString c (ResourceT IO) [PsResult]
consumer)
  where
    consumer :: ConduitM ByteString c (ResourceT IO) [PsResult]
consumer = ConduitT ByteString ByteString (ResourceT IO) ()
forall (m :: * -> *).
Monad m =>
ConduitT ByteString ByteString m ()
CB.lines
            ConduitT ByteString ByteString (ResourceT IO) ()
-> ConduitM ByteString c (ResourceT IO) [PsResult]
-> ConduitM ByteString c (ResourceT IO) [PsResult]
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| (ByteString -> [ByteString])
-> ConduitT ByteString [ByteString] (ResourceT IO) ()
forall (m :: * -> *) a b. Monad m => (a -> b) -> ConduitT a b m ()
CL.map ByteString -> [ByteString]
BS.words
            ConduitT ByteString [ByteString] (ResourceT IO) ()
-> ConduitM [ByteString] c (ResourceT IO) [PsResult]
-> ConduitM ByteString c (ResourceT IO) [PsResult]
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| ([ByteString] -> PsResult)
-> ConduitT [ByteString] PsResult (ResourceT IO) ()
forall (m :: * -> *) a b. Monad m => (a -> b) -> ConduitT a b m ()
CL.map [ByteString] -> PsResult
toPsResult
            ConduitT [ByteString] PsResult (ResourceT IO) ()
-> ConduitM PsResult c (ResourceT IO) [PsResult]
-> ConduitM [ByteString] c (ResourceT IO) [PsResult]
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| (PsResult -> Bool) -> ConduitT PsResult PsResult (ResourceT IO) ()
forall (m :: * -> *) a. Monad m => (a -> Bool) -> ConduitT a a m ()
CL.filter PsResult -> Bool
mighty
            ConduitT PsResult PsResult (ResourceT IO) ()
-> ConduitM PsResult c (ResourceT IO) [PsResult]
-> ConduitM PsResult c (ResourceT IO) [PsResult]
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| ConduitM PsResult c (ResourceT IO) [PsResult]
forall (m :: * -> *) a o. Monad m => ConduitT a o m [a]
CL.consume
    commandName :: PsResult -> ByteString
commandName = [ByteString] -> ByteString
forall a. [a] -> a
last ([ByteString] -> ByteString)
-> (PsResult -> [ByteString]) -> PsResult -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> ByteString -> [ByteString]
split Char
'/' (ByteString -> [ByteString])
-> (PsResult -> ByteString) -> PsResult -> [ByteString]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PsResult -> ByteString
command
    mighty :: PsResult -> Bool
mighty PsResult
ps = ByteString
"mighty" ByteString -> ByteString -> Bool
`BS.isInfixOf` ByteString
name
             Bool -> Bool -> Bool
&& Bool -> Bool
not (ByteString
"mightyctl" ByteString -> ByteString -> Bool
`BS.isInfixOf` ByteString
name)
        where
          name :: ByteString
name = PsResult -> ByteString
commandName PsResult
ps

----------------------------------------------------------------

findParent :: [PsResult] -> [PsResult]
findParent :: [PsResult] -> [PsResult]
findParent [PsResult]
ps = [PsResult] -> [PsResult]
deleteAloneChild ([PsResult] -> [PsResult]) -> [PsResult] -> [PsResult]
forall a b. (a -> b) -> a -> b
$ [PsResult]
masters [PsResult] -> [PsResult] -> [PsResult]
forall a. [a] -> [a] -> [a]
++ [PsResult]
candidates
  where
    iAmMaster :: PsResult -> Bool
iAmMaster PsResult
p = PsResult -> ProcessID
ppid PsResult
p ProcessID -> ProcessID -> Bool
forall a. Eq a => a -> a -> Bool
== ProcessID
1
    masters :: [PsResult]
masters = (PsResult -> Bool) -> [PsResult] -> [PsResult]
forall a. (a -> Bool) -> [a] -> [a]
filter PsResult -> Bool
iAmMaster [PsResult]
ps
    rest :: [PsResult]
rest    = (PsResult -> Bool) -> [PsResult] -> [PsResult]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not(Bool -> Bool) -> (PsResult -> Bool) -> PsResult -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
.PsResult -> Bool
iAmMaster) [PsResult]
ps
    candidates :: [PsResult]
candidates = ([PsResult] -> PsResult) -> [[PsResult]] -> [PsResult]
forall a b. (a -> b) -> [a] -> [b]
map [PsResult] -> PsResult
forall a. [a] -> a
head
               ([[PsResult]] -> [PsResult]) -> [[PsResult]] -> [PsResult]
forall a b. (a -> b) -> a -> b
$ ([PsResult] -> Bool) -> [[PsResult]] -> [[PsResult]]
forall a. (a -> Bool) -> [a] -> [a]
filter (\[PsResult]
xs -> [PsResult] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [PsResult]
xs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1) -- master is alone
               ([[PsResult]] -> [[PsResult]]) -> [[PsResult]] -> [[PsResult]]
forall a b. (a -> b) -> a -> b
$ (PsResult -> PsResult -> Bool) -> [PsResult] -> [[PsResult]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy (ProcessID -> ProcessID -> Bool
forall a. Eq a => a -> a -> Bool
(==) (ProcessID -> ProcessID -> Bool)
-> (PsResult -> ProcessID) -> PsResult -> PsResult -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` PsResult -> ProcessID
ppid)
               ([PsResult] -> [[PsResult]]) -> [PsResult] -> [[PsResult]]
forall a b. (a -> b) -> a -> b
$ (PsResult -> ProcessID) -> [PsResult] -> [PsResult]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn PsResult -> ProcessID
ppid [PsResult]
rest


deleteAloneChild :: [PsResult] -> [PsResult]
deleteAloneChild :: [PsResult] -> [PsResult]
deleteAloneChild [] = []
deleteAloneChild (PsResult
p:[PsResult]
ps) = PsResult
p PsResult -> [PsResult] -> [PsResult]
forall a. a -> [a] -> [a]
: [PsResult] -> [PsResult]
deleteAloneChild ((PsResult -> Bool) -> [PsResult] -> [PsResult]
forall a. (a -> Bool) -> [a] -> [a]
filter PsResult -> Bool
noParent [PsResult]
ps)
  where
    parent :: ProcessID
parent = PsResult -> ProcessID
pid PsResult
p
    noParent :: PsResult -> Bool
noParent PsResult
x = PsResult -> ProcessID
ppid PsResult
x ProcessID -> ProcessID -> Bool
forall a. Eq a => a -> a -> Bool
/= ProcessID
parent

----------------------------------------------------------------

-- | Getting the process id of a running Mighty.
getMightyPid :: IO [ProcessID]
getMightyPid :: IO [ProcessID]
getMightyPid = (PsResult -> ProcessID) -> [PsResult] -> [ProcessID]
forall a b. (a -> b) -> [a] -> [b]
map PsResult -> ProcessID
pid ([PsResult] -> [ProcessID])
-> ([PsResult] -> [PsResult]) -> [PsResult] -> [ProcessID]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [PsResult] -> [PsResult]
findParent ([PsResult] -> [ProcessID]) -> IO [PsResult] -> IO [ProcessID]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO [PsResult]
runPS

----------------------------------------------------------------

split :: Char -> ByteString -> [ByteString]
split :: Char -> ByteString -> [ByteString]
split Char
_ ByteString
"" = []
split Char
c ByteString
s = case (Char -> Bool) -> ByteString -> (ByteString, ByteString)
BS.break (Char
cChar -> Char -> Bool
forall a. Eq a => a -> a -> Bool
==) ByteString
s of
    (ByteString
"",ByteString
r)  -> Char -> ByteString -> [ByteString]
split Char
c (ByteString -> ByteString
BS.tail ByteString
r)
    (ByteString
s',ByteString
"") -> [ByteString
s']
    (ByteString
s',ByteString
r)  -> ByteString
s' ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: Char -> ByteString -> [ByteString]
split Char
c (ByteString -> ByteString
BS.tail ByteString
r)