{-# LANGUAGE OverloadedStrings #-}
module ProcessUtils
    ( ExitCode(ExitFailure,ExitSuccess)
    , systemEnv
    , readPipe
    , readProcessWithErrorH
    , runExternal
    , SpawnError(..)
    , spawnDetached
    ) where

import GHC.IO.Exception ( ioException, IOErrorType(..) )
import System.FilePath
import System.Directory
import System.Process
import System.Posix.Signals
import System.Posix.Process
import System.Process.Internals (runGenProcess_,defaultSignal)
import System.Environment
import Data.Maybe ( isNothing, fromMaybe, isJust, fromJust )
import System.IO.Error ( mkIOError, ioeSetErrorString ) 
import System.Exit ( ExitCode(..) )
import System.IO
import Control.Applicative
import Control.Exception (bracket)
import qualified Data.ByteString as S
import ControlMaybe (handleIO,handleIO_)
import System.Posix.Files (fileAccess)
import System.Posix.User (getRealUserID,getEffectiveUserID)
import DotLock (dotlock_disable)

-- | systemEnv
-- This is like System.Process.system except that it lets you set
-- some environment variables.
systemEnv :: [(String, String)] -> String -> IO ExitCode
systemEnv _    ""  = 
    ioException (ioeSetErrorString (mkIOError InvalidArgument "system" Nothing Nothing) "null command")
systemEnv vars cmd = do
    env0 <- getEnvironment
    let env1 = filter (isNothing . flip lookup vars . fst) env0
        env = vars ++ env1
    syncProcess "system" $ (shell cmd) {env=Just env}
 where
    -- This is a non-exported function from System.Process
    syncProcess fun c = do
      -- The POSIX version of system needs to do some manipulation of signal
      -- handlers.  Since we're going to be synchronously waiting for the child,
      -- we want to ignore ^C in the parent, but handle it the default way
      -- in the child (using SIG_DFL isn't really correct, it should be the
      -- original signal handler, but the GHC RTS will have already set up
      -- its own handler and we don't want to use that).
      old_int  <- installHandler sigINT  Ignore Nothing
      old_quit <- installHandler sigQUIT Ignore Nothing
      (_,_,_,p) <- runGenProcess_ fun c
            (Just defaultSignal) (Just defaultSignal)
      r <- waitForProcess p
      _ <- installHandler sigINT  old_int Nothing
      _ <- installHandler sigQUIT old_quit Nothing
      return r

-- | readPipe
--
-- This is like System.Process.readProcess but instead of a single
-- command with arguments, you pass it a list which it then executes
-- piping the the output from the prior command into the input for
-- the next command. Like readProcess the next parameter is a string
-- which is sent as input into the initial standard in.
--
--  > readPipe [("grep",["HostConfig"]),("wc",["-l"]) =<< readFile "input"
--
-- is equivalent to shell code:
--
--  > grep HostConfig input | wc -l
--
readPipe :: [(FilePath,[String])] -> String -> IO String
readPipe [(cmd,args)] input = readProcess cmd args input
readPipe ((cmd,args):xs) stdin0 = do
    let p = (shell cmd) { std_in = CreatePipe
                        , std_out = CreatePipe
                        , std_err = CreatePipe
                        , cmdspec = RawCommand cmd args
                        }
    (Just sinh,Just sout,Just serr, ph) <- createProcess p
    hPutStr sinh stdin0
    hClose sinh
    readPipe0 xs sout
    where
        readPipe0 :: [(FilePath,[String])] -> Handle -> IO String
        readPipe0 ((cmd,args):xs) stdin0 = do
            let p = (shell cmd)
                        { std_in = UseHandle stdin0
                        , std_out = CreatePipe
                        , std_err = CreatePipe
                        , cmdspec = RawCommand cmd args
                        }
            (Nothing,Just sout,Just serr, ph) <- createProcess p
            readPipe0 xs sout
        readPipe0 [] h = hGetContents h


readProcessWithErrorH :: FilePath -> [String] -> String -> Handle -> IO String
readProcessWithErrorH cmd args stdin erH = do
    let p = (shell cmd) { std_in = CreatePipe
                        , std_out = CreatePipe
                        , std_err = UseHandle erH
                        , cmdspec = RawCommand cmd args
                        }
    bracket (createProcess p) 
            (\(Just sinh,Just sout,_, ph) ->
                mapM_ hClose [sinh,sout]         )
            (\(Just sinh,Just sout,_, ph) -> do
                hPutStr sinh stdin
                hClose sinh
                hGetContents sout                )

-- | run an external command.
--
-- Arguments:
--
--   [ cmd ]   command to run
--
--   [ input ] input to command (if any)
--
-- Returns:
--
--   [ code ]   exit code returned by command
--
--   [ output ] captured stdout of command
--
--   [ err ]    captured stderr of command
runExternal :: String -> Maybe String -> IO (Int,(S.ByteString,S.ByteString))
runExternal cmd input = do
    cwd <- getCurrentDirectory
    {-
    putStr $ takeFileName cwd ++ "> " ++ cmd
                    ++ case input of
                            Nothing -> "\n"
                            Just s  -> " <<EOF\n" ++ s ++ "EOF\n"
    -}
    let p = (shell cmd) { std_in  = maybe Inherit (const CreatePipe) input
                        , std_out = CreatePipe
                        , std_err = CreatePipe }
    (mb_in,mb_out,mb_err,process) <- createProcess p
    case mb_in of
        Nothing -> return ()
        Just hin -> hPutStr hin (fromMaybe "" input) >> hClose hin
    out <- case mb_out of
            Just hout -> S.hGetContents hout
            Nothing   -> return ""
    err <- case mb_err of
            Just herr -> S.hGetContents herr
            Nothing   -> return ""
    code <- do
        c <- waitForProcess process
        return $ case c of
                    ExitSuccess   -> 0
                    ExitFailure x -> x
    return (code,(out,err))


data SpawnError = SpawnOK | SpawnUIDCheck | SpawnBadAccess IOError | SpawnWaitPID | SpawnFork
 deriving (Eq,Show)

spawnDetached :: FilePath -> [String] -> Maybe [(String,String)] -> IO SpawnError
spawnDetached filepath args env = do
    uidcheck <- (==) <$> getRealUserID <*> getEffectiveUserID
    if not uidcheck then return SpawnUIDCheck else do
    accesscheck <- do
        handleIO (return . Just . SpawnBadAccess) $ do
            b <- fileAccess filepath False False True
            if not b then return $ Just $ SpawnBadAccess (userError $ "No execute permission on "++filepath)
                     else return Nothing
    if isJust accesscheck then return (fromJust accesscheck) else do
    handleIO_ (return SpawnFork) $ do
        pid <- forkProcess $ do
            dotlock_disable -- Disable dotlock atexit handler.
                            -- This wouldn't be neccessary if we used exec*
                            -- instead of fork.
            -- clean up secure memory:
            -- gcry_control (GCRYCTL_TERM_SECMEM);
            -- if (setsid() == -1 || chdir ("/")) _exit (1);

            -- Double fork to let init take over the new child.
            handleIO_ (exitImmediately $ ExitFailure 1) $ do
                pid2 <- forkProcess $ executeFile filepath False args env
                exitImmediately ExitSuccess
            exitImmediately ExitSuccess
        handleIO_ (return SpawnWaitPID) $ do
            _ <- getProcessStatus True False pid
            return SpawnOK