{-# 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
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
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 {
    uid :: ByteString
uid     = ByteString
a
  , pid :: ProcessID
pid     = forall b a. b -> (a -> b) -> Maybe a -> b
maybe ProcessID
0 (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) forall a b. (a -> b) -> a -> b
$ ByteString -> Maybe (Int, ByteString)
BS.readInt ByteString
b
  , ppid :: ProcessID
ppid    = forall b a. b -> (a -> b) -> Maybe a -> b
maybe ProcessID
0 (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) 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 = forall a b. (a, b) -> b
snd forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a. MonadUnliftIO m => ResourceT m a -> m a
runResourceT (forall (m :: * -> *) a.
MonadIO m =>
String -> ConduitT ByteString Void m a -> m (ExitCode, a)
sourceCmdWithConsumer String
"ps -ef" forall {c}. ConduitT ByteString c (ResourceT IO) [PsResult]
consumer)
  where
    consumer :: ConduitT ByteString c (ResourceT IO) [PsResult]
consumer = forall (m :: * -> *).
Monad m =>
ConduitT ByteString ByteString m ()
CB.lines
            forall (m :: * -> *) a b c r.
Monad m =>
ConduitT a b m () -> ConduitT b c m r -> ConduitT a c m r
.| forall (m :: * -> *) a b. Monad m => (a -> b) -> ConduitT a b m ()
CL.map ByteString -> [ByteString]
BS.words
            forall (m :: * -> *) a b c r.
Monad m =>
ConduitT a b m () -> ConduitT b c m r -> ConduitT a c m r
.| forall (m :: * -> *) a b. Monad m => (a -> b) -> ConduitT a b m ()
CL.map [ByteString] -> PsResult
toPsResult
            forall (m :: * -> *) a b c r.
Monad m =>
ConduitT a b m () -> ConduitT b c m r -> ConduitT a c m r
.| forall (m :: * -> *) a. Monad m => (a -> Bool) -> ConduitT a a m ()
CL.filter PsResult -> Bool
mighty
            forall (m :: * -> *) a b c r.
Monad m =>
ConduitT a b m () -> ConduitT b c m r -> ConduitT a c m r
.| forall (m :: * -> *) a o. Monad m => ConduitT a o m [a]
CL.consume
    commandName :: PsResult -> ByteString
commandName = forall a. [a] -> a
last forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> ByteString -> [ByteString]
split Char
'/' 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 forall a b. (a -> b) -> a -> b
$ [PsResult]
masters forall a. [a] -> [a] -> [a]
++ [PsResult]
candidates
  where
    iAmMaster :: PsResult -> Bool
iAmMaster PsResult
p = PsResult -> ProcessID
ppid PsResult
p forall a. Eq a => a -> a -> Bool
== ProcessID
1
    masters :: [PsResult]
masters = forall a. (a -> Bool) -> [a] -> [a]
filter PsResult -> Bool
iAmMaster [PsResult]
ps
    rest :: [PsResult]
rest    = forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
notforall b c a. (b -> c) -> (a -> b) -> a -> c
.PsResult -> Bool
iAmMaster) [PsResult]
ps
    candidates :: [PsResult]
candidates = forall a b. (a -> b) -> [a] -> [b]
map forall a. [a] -> a
head
               forall a b. (a -> b) -> a -> b
$ forall a. (a -> Bool) -> [a] -> [a]
filter (\[PsResult]
xs -> forall (t :: * -> *) a. Foldable t => t a -> Int
length [PsResult]
xs forall a. Eq a => a -> a -> Bool
== Int
1) -- master is alone
               forall a b. (a -> b) -> a -> b
$ forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy (forall a. Eq a => a -> a -> Bool
(==) forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` PsResult -> ProcessID
ppid)
               forall a b. (a -> b) -> a -> b
$ 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 forall a. a -> [a] -> [a]
: [PsResult] -> [PsResult]
deleteAloneChild (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 forall a. Eq a => a -> a -> Bool
/= ProcessID
parent

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

-- | Getting the process id of a running Mighty.
getMightyPid :: IO [ProcessID]
getMightyPid :: IO [ProcessID]
getMightyPid = forall a b. (a -> b) -> [a] -> [b]
map PsResult -> ProcessID
pid forall b c a. (b -> c) -> (a -> b) -> a -> c
. [PsResult] -> [PsResult]
findParent 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
cforall a. Eq a => a -> a -> Bool
==) ByteString
s of
    (ByteString
"",ByteString
r)  -> Char -> ByteString -> [ByteString]
split Char
c (HasCallStack => ByteString -> ByteString
BS.tail ByteString
r)
    (ByteString
s',ByteString
"") -> [ByteString
s']
    (ByteString
s',ByteString
r)  -> ByteString
s' forall a. a -> [a] -> [a]
: Char -> ByteString -> [ByteString]
split Char
c (HasCallStack => ByteString -> ByteString
BS.tail ByteString
r)